From d51af8a5835347d426d4bece36d36f7586888c37 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 13 Dec 2021 13:28:40 -0500 Subject: [PATCH] Agent: Add IPScanResults dataclass --- monkey/infection_monkey/master/__init__.py | 1 + .../master/ip_scan_results.py | 14 ++++++++++ monkey/infection_monkey/master/ip_scanner.py | 22 ++++++---------- monkey/infection_monkey/master/propagator.py | 18 +++++-------- .../master/test_ip_scanner.py | 26 +++++++++++-------- .../master/test_propagator.py | 14 +++++----- 6 files changed, 52 insertions(+), 43 deletions(-) create mode 100644 monkey/infection_monkey/master/ip_scan_results.py diff --git a/monkey/infection_monkey/master/__init__.py b/monkey/infection_monkey/master/__init__.py index 21ef8f9b6..fda536194 100644 --- a/monkey/infection_monkey/master/__init__.py +++ b/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 diff --git a/monkey/infection_monkey/master/ip_scan_results.py b/monkey/infection_monkey/master/ip_scan_results.py new file mode 100644 index 000000000..98f7b6646 --- /dev/null +++ b/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] diff --git a/monkey/infection_monkey/master/ip_scanner.py b/monkey/infection_monkey/master/ip_scanner.py index 62bf9c7d8..26a321212 100644 --- a/monkey/infection_monkey/master/ip_scanner.py +++ b/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: diff --git a/monkey/infection_monkey/master/propagator.py b/monkey/infection_monkey/master/propagator.py index 0d63bc904..1d6e4462e 100644 --- a/monkey/infection_monkey/master/propagator.py +++ b/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) diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py b/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py index 12e822fa3..22c850837 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py +++ b/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): diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_propagator.py b/monkey/tests/unit_tests/infection_monkey/master/test_propagator.py index cec779aa5..d8f65b54e 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_propagator.py +++ b/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):