diff --git a/monkey/infection_monkey/i_puppet.py b/monkey/infection_monkey/i_puppet.py index 518b299b6..285e32bca 100644 --- a/monkey/infection_monkey/i_puppet.py +++ b/monkey/infection_monkey/i_puppet.py @@ -58,11 +58,20 @@ class IPuppet(metaclass=abc.ABCMeta): """ @abc.abstractmethod - def fingerprint(self, name: str, host: str) -> FingerprintData: + def fingerprint( + self, + name: str, + host: str, + ping_scan_data: PingScanData, + port_scan_data: Dict[int, PortScanData], + ) -> FingerprintData: """ Runs a fingerprinter against a remote host :param str name: The name of the fingerprinter to run :param str host: The domain name or IP address of a host + :param PingScanData ping_scan_data: Data retrieved from the target host via ICMP + :param Dict[int, PortScanData] port_scan_data: Data retrieved from the target host via a TCP + port scan :return: The data collected by running the fingerprinter on the specified host :rtype: FingerprintData """ diff --git a/monkey/infection_monkey/master/ip_scanner.py b/monkey/infection_monkey/master/ip_scanner.py index 26a321212..cf77ea54d 100644 --- a/monkey/infection_monkey/master/ip_scanner.py +++ b/monkey/infection_monkey/master/ip_scanner.py @@ -5,7 +5,13 @@ from queue import Queue from threading import Event from typing import Callable, Dict, List -from infection_monkey.i_puppet import FingerprintData, IPuppet, PortScanData, PortStatus +from infection_monkey.i_puppet import ( + FingerprintData, + IPuppet, + PingScanData, + PortScanData, + PortStatus, +) from . import IPScanResults from .threading_utils import create_daemon_thread @@ -56,7 +62,9 @@ class IPScanner: fingerprint_data = {} if IPScanner._found_open_port(port_scan_data): fingerprinters = options["fingerprinters"] - fingerprint_data = self._run_fingerprinters(ip, fingerprinters, stop) + fingerprint_data = self._run_fingerprinters( + ip, fingerprinters, ping_scan_data, port_scan_data, stop + ) scan_results = IPScanResults(ping_scan_data, port_scan_data, fingerprint_data) results_callback(ip, scan_results) @@ -92,7 +100,12 @@ class IPScanner: return False def _run_fingerprinters( - self, ip: str, fingerprinters: List[str], stop: Event + self, + ip: str, + fingerprinters: List[str], + ping_scan_data: PingScanData, + port_scan_data: Dict[int, PortScanData], + stop: Event, ) -> Dict[str, FingerprintData]: fingerprint_data = {} @@ -100,6 +113,6 @@ class IPScanner: if stop.is_set(): break - fingerprint_data[f] = self._puppet.fingerprint(f, ip) + fingerprint_data[f] = self._puppet.fingerprint(f, ip, ping_scan_data, port_scan_data) return fingerprint_data diff --git a/monkey/infection_monkey/master/mock_master.py b/monkey/infection_monkey/master/mock_master.py index 551ff886c..3844ef590 100644 --- a/monkey/infection_monkey/master/mock_master.py +++ b/monkey/infection_monkey/master/mock_master.py @@ -88,13 +88,13 @@ class MockMaster(IMaster): machine_1 = self._hosts["10.0.0.1"] machine_3 = self._hosts["10.0.0.3"] - self._puppet.fingerprint("SMBFinger", machine_1) + self._puppet.fingerprint("SMBFinger", machine_1, None, None) self._telemetry_messenger.send_telemetry(ScanTelem(machine_1)) - self._puppet.fingerprint("SMBFinger", machine_3) + self._puppet.fingerprint("SMBFinger", machine_3, None, None) self._telemetry_messenger.send_telemetry(ScanTelem(machine_3)) - self._puppet.fingerprint("HTTPFinger", machine_3) + self._puppet.fingerprint("HTTPFinger", machine_3, None, None) self._telemetry_messenger.send_telemetry(ScanTelem(machine_3)) logger.info("Finished running fingerprinters on potential victims") diff --git a/monkey/infection_monkey/puppet/mock_puppet.py b/monkey/infection_monkey/puppet/mock_puppet.py index f8c76d843..d5c8fa2f8 100644 --- a/monkey/infection_monkey/puppet/mock_puppet.py +++ b/monkey/infection_monkey/puppet/mock_puppet.py @@ -194,7 +194,13 @@ class MockPuppet(IPuppet): return _get_empty_results(port) - def fingerprint(self, name: str, host: str) -> FingerprintData: + def fingerprint( + self, + name: str, + host: str, + ping_scan_data: PingScanData, + port_scan_data: Dict[int, PortScanData], + ) -> FingerprintData: logger.debug(f"fingerprint({name}, {host})") empty_fingerprint_data = FingerprintData(None, None, {}) 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 22c850837..3b071eb9a 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 @@ -229,7 +229,7 @@ def test_interrupt_port_scanning(callback, scan_config, stop): def test_interrupt_fingerprinting(callback, scan_config, stop): - def stopable_fingerprint(port, *_): + def stopable_fingerprint(*_): # Block all threads here until 2 threads reach this barrier, then set stop # and test that neither thread scans any more ports stopable_fingerprint.barrier.wait()