Per-host speed limit (#1598)

Author: kazet

artemis/config.py

@@ -1,5 +1,5 @@
 import os
-from typing import Annotated, Any, List, get_type_hints
+from typing import Annotated, Any, List, Optional, get_type_hints
 
 import decouple
 
@@ -166,6 +166,11 @@ class Limits:
             "Default request timeout (for all protocols).",
         ] = get_config("REQUEST_TIMEOUT_SECONDS", default=5, cast=int)
 
+        SCAN_SPEED_OVERRIDES_FILE: Annotated[
+            Optional[str],
+            "A JSON file with a dictionary mapping from IP to scan speed - use if you want to slow down scanning of particular hosts.",
+        ] = get_config("SCAN_SPEED_OVERRIDES_FILE", default="", cast=str)
+
         REQUESTS_PER_SECOND: Annotated[
             float,
             """

artemis/module_base.py

@@ -1,5 +1,6 @@
 import datetime
 import fcntl
+import json
 import logging
 import random
 import shutil
@@ -8,7 +9,7 @@
 import traceback
 import urllib.parse
 from ipaddress import ip_address
-from typing import Any, Callable, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import timeout_decorator
 from karton.core import Karton, Task
@@ -89,6 +90,12 @@ def __init__(self, db: Optional[DB] = None, *args, **kwargs) -> None:  # type: i
                     hosts_file.write(additional_data_file.read())
                     fcntl.flock(hosts_file.fileno(), fcntl.LOCK_UN)
 
+        if Config.Limits.SCAN_SPEED_OVERRIDES_FILE:
+            with open(Config.Limits.SCAN_SPEED_OVERRIDES_FILE, "r") as f:
+                self._scan_speed_overrides: Dict[str, float] = json.load(f)
+        else:
+            self._scan_speed_overrides = {}
+
         if Config.Miscellaneous.BLOCKLIST_FILE:
             self._blocklist = load_blocklist(Config.Miscellaneous.BLOCKLIST_FILE)
         else:
@@ -288,12 +295,17 @@ def _single_iteration(self) -> int:
             if "requests_per_second_override" in task.payload_persistent
         ]
 
+        for task in tasks:
+            destination = self._get_scan_destination(task)
+            if destination in self._scan_speed_overrides:
+                requests_per_second_overrides.append(self._scan_speed_overrides[destination])
+
         self.requests_per_second_for_current_tasks = min(  # type: ignore
             requests_per_second_overrides if requests_per_second_overrides else [Config.Limits.REQUESTS_PER_SECOND]
         )
 
         if requests_per_second_overrides:
-            self.log.info("Overriding requests per second to %f", self.requests_per_second_for_current_tasks)
+            self.log.info("Setting requests per second to %f", self.requests_per_second_for_current_tasks)
 
         if len(tasks):
             time_start = time.time()

artemis/modules/port_scanner.py

@@ -125,7 +125,7 @@ def _scan(self, target_ips: List[str]) -> Dict[str, Dict[str, Dict[str, Any]]]:
                     str(Config.Modules.PortScanner.PORT_SCANNER_TIMEOUT_MILLISECONDS),
                 ]
                 + (
-                    ["-rate", str(int(self.requests_per_second_for_current_tasks) * len(new_target_ips))]
+                    ["-rate", str(max(1, int(self.requests_per_second_for_current_tasks)) * len(new_target_ips))]
                     if int(self.requests_per_second_for_current_tasks)
                     else []
                 ),