Agent: Add IPScanResults dataclass

monkey/infection_monkey/master/__init__.py

@@ -1,3 +1,4 @@
+from .ip_scan_results import IPScanResults
 from .ip_scanner import IPScanner
 from .propagator import Propagator
 from .automated_master import AutomatedMaster

monkey/infection_monkey/master/ip_scan_results.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+from typing import Dict
+
+from infection_monkey.i_puppet import FingerprintData, PingScanData, PortScanData
+
+Port = int
+FingerprinterName = str
+
+
+@dataclass
+class IPScanResults:
+    ping_scan_data: PingScanData
+    port_scan_data: Dict[Port, PortScanData]
+    fingerprint_data: Dict[FingerprinterName, FingerprintData]

monkey/infection_monkey/master/ip_scanner.py

@@ -5,24 +5,15 @@ from queue import Queue
 from threading import Event
 from typing import Callable, Dict, List
 
-from infection_monkey.i_puppet import (
-    FingerprintData,
-    IPuppet,
-    PingScanData,
-    PortScanData,
-    PortStatus,
-)
+from infection_monkey.i_puppet import FingerprintData, IPuppet, PortScanData, PortStatus
 
+from . import IPScanResults
 from .threading_utils import create_daemon_thread
 
 logger = logging.getLogger()
 
 IP = str
-Port = int
-FingerprinterName = str
-Callback = Callable[
-    [IP, PingScanData, Dict[Port, PortScanData], Dict[FingerprinterName, FingerprintData]], None
-]
+Callback = Callable[[IP, IPScanResults], None]
 
 
 class IPScanner:
@@ -67,7 +58,8 @@ class IPScanner:
                     fingerprinters = options["fingerprinters"]
                     fingerprint_data = self._run_fingerprinters(ip, fingerprinters, stop)
 
-                results_callback(ip, ping_scan_data, port_scan_data, fingerprint_data)
+                scan_results = IPScanResults(ping_scan_data, port_scan_data, fingerprint_data)
+                results_callback(ip, scan_results)
 
             logger.debug(
                 f"Detected the stop signal, scanning thread {threading.get_ident()} exiting"
@@ -99,7 +91,9 @@ class IPScanner:
 
         return False
 
-    def _run_fingerprinters(self, ip: str, fingerprinters: List[str], stop: Event):
+    def _run_fingerprinters(
+        self, ip: str, fingerprinters: List[str], stop: Event
+    ) -> Dict[str, FingerprintData]:
         fingerprint_data = {}
 
         for f in fingerprinters:

monkey/infection_monkey/master/propagator.py

@@ -8,7 +8,7 @@ from infection_monkey.model.host import VictimHost
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.telemetry.scan_telem import ScanTelem
 
-from . import IPScanner
+from . import IPScanner, IPScanResults
 from .threading_utils import create_daemon_thread
 
 logger = logging.getLogger()
@@ -51,18 +51,14 @@ class Propagator:
 
         logger.info("Finished network scan")
 
-    def _process_scan_results(
-        self,
-        ip: str,
-        ping_scan_data: PingScanData,
-        port_scan_data: Dict[int, PortScanData],
-        fingerprint_data: Dict[str, FingerprintData],
-    ):
+    def _process_scan_results(self, ip: str, scan_results: IPScanResults):
         victim_host = VictimHost(ip)
 
-        Propagator._process_ping_scan_results(victim_host, ping_scan_data)
-        has_open_port = Propagator._process_tcp_scan_results(victim_host, port_scan_data)
-        Propagator._process_fingerprinter_results(victim_host, fingerprint_data)
+        Propagator._process_ping_scan_results(victim_host, scan_results.ping_scan_data)
+        has_open_port = Propagator._process_tcp_scan_results(
+            victim_host, scan_results.port_scan_data
+        )
+        Propagator._process_fingerprinter_results(victim_host, scan_results.fingerprint_data)
 
         if has_open_port:
             self._hosts_to_exploit.put(victim_host)

monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py

@@ -51,7 +51,11 @@ def assert_port_status(port_scan_data, expected_open_ports: Set[int]):
             assert psd.status == PortStatus.CLOSED
 
 
-def assert_scan_results(ip, ping_scan_data, port_scan_data, fingerprint_data):
+def assert_scan_results(ip, scan_results):
+    ping_scan_data = scan_results.ping_scan_data
+    port_scan_data = scan_results.port_scan_data
+    fingerprint_data = scan_results.fingerprint_data
+
     if ip == "10.0.0.1":
         assert_scan_results_no_1(ping_scan_data, port_scan_data, fingerprint_data)
     elif ip == "10.0.0.3":
@@ -149,8 +153,8 @@ def test_scan_single_ip(callback, scan_config, stop):
 
     callback.assert_called_once()
 
-    (ip, ping_scan_data, port_scan_data, fingerprint_data) = callback.call_args_list[0][0]
-    assert_scan_results(ip, ping_scan_data, port_scan_data, fingerprint_data)
+    (ip, scan_results) = callback.call_args_list[0][0]
+    assert_scan_results(ip, scan_results)
 
 
 def test_scan_multiple_ips(callback, scan_config, stop):
@@ -161,17 +165,17 @@ def test_scan_multiple_ips(callback, scan_config, stop):
 
     assert callback.call_count == 4
 
-    (ip, ping_scan_data, port_scan_data, fingerprint_data) = callback.call_args_list[0][0]
-    assert_scan_results(ip, ping_scan_data, port_scan_data, fingerprint_data)
+    (ip, scan_results) = callback.call_args_list[0][0]
+    assert_scan_results(ip, scan_results)
 
-    (ip, ping_scan_data, port_scan_data, fingerprint_data) = callback.call_args_list[1][0]
-    assert_scan_results(ip, ping_scan_data, port_scan_data, fingerprint_data)
+    (ip, scan_results) = callback.call_args_list[1][0]
+    assert_scan_results(ip, scan_results)
 
-    (ip, ping_scan_data, port_scan_data, fingerprint_data) = callback.call_args_list[2][0]
-    assert_scan_results(ip, ping_scan_data, port_scan_data, fingerprint_data)
+    (ip, scan_results) = callback.call_args_list[2][0]
+    assert_scan_results(ip, scan_results)
 
-    (ip, ping_scan_data, port_scan_data, fingerprint_data) = callback.call_args_list[3][0]
-    assert_scan_results(ip, ping_scan_data, port_scan_data, fingerprint_data)
+    (ip, scan_results) = callback.call_args_list[3][0]
+    assert_scan_results(ip, scan_results)
 
 
 def test_scan_lots_of_ips(callback, scan_config, stop):

monkey/tests/unit_tests/infection_monkey/master/test_propagator.py

@@ -1,11 +1,11 @@
 from threading import Event
 
 from infection_monkey.i_puppet import FingerprintData, PingScanData, PortScanData, PortStatus
-from infection_monkey.master import Propagator
+from infection_monkey.master import IPScanResults, Propagator
 
 empty_fingerprint_data = FingerprintData(None, None, {})
 
-dot_1_results = (
+dot_1_results = IPScanResults(
     PingScanData(True, "windows"),
     {
         22: PortScanData(22, PortStatus.CLOSED, None, None),
@@ -19,7 +19,7 @@ dot_1_results = (
     },
 )
 
-dot_3_results = (
+dot_3_results = IPScanResults(
     PingScanData(True, "linux"),
     {
         22: PortScanData(22, PortStatus.OPEN, "SSH BANNER", "tcp-22"),
@@ -42,7 +42,7 @@ dot_3_results = (
     },
 )
 
-dead_host_results = (
+dead_host_results = IPScanResults(
     PingScanData(False, None),
     {
         22: PortScanData(22, PortStatus.CLOSED, None, None),
@@ -79,11 +79,11 @@ class MockIPScanner:
     def scan(self, ips_to_scan, _, results_callback, stop):
         for ip in ips_to_scan:
             if ip.endswith(".1"):
-                results_callback(ip, *dot_1_results)
+                results_callback(ip, dot_1_results)
             elif ip.endswith(".3"):
-                results_callback(ip, *dot_3_results)
+                results_callback(ip, dot_3_results)
             else:
-                results_callback(ip, *dead_host_results)
+                results_callback(ip, dead_host_results)
 
 
 def test_scan_result_processing(telemetry_messenger_spy):