diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 784046323..57b8f52b2 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -6,6 +6,7 @@ from typing import Any, Callable, Dict, List, Tuple from infection_monkey.i_control_channel import IControlChannel from infection_monkey.i_master import IMaster from infection_monkey.i_puppet import IPuppet +from infection_monkey.model import VictimHostFactory from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger from infection_monkey.telemetry.post_breach_telem import PostBreachTelem from infection_monkey.telemetry.system_info_telem import SystemInfoTelem @@ -27,6 +28,7 @@ class AutomatedMaster(IMaster): self, puppet: IPuppet, telemetry_messenger: ITelemetryMessenger, + victim_host_factory: VictimHostFactory, control_channel: IControlChannel, ): self._puppet = puppet @@ -34,7 +36,7 @@ class AutomatedMaster(IMaster): self._control_channel = control_channel ip_scanner = IPScanner(self._puppet, NUM_SCAN_THREADS) - self._propagator = Propagator(self._telemetry_messenger, ip_scanner) + self._propagator = Propagator(self._telemetry_messenger, ip_scanner, victim_host_factory) self._stop = threading.Event() self._master_thread = create_daemon_thread(target=self._run_master_thread) diff --git a/monkey/infection_monkey/master/propagator.py b/monkey/infection_monkey/master/propagator.py index 916297110..78e08a98d 100644 --- a/monkey/infection_monkey/master/propagator.py +++ b/monkey/infection_monkey/master/propagator.py @@ -4,7 +4,7 @@ from threading import Event, Thread from typing import Dict from infection_monkey.i_puppet import FingerprintData, PingScanData, PortScanData, PortStatus -from infection_monkey.model.host import VictimHost +from infection_monkey.model import VictimHost, VictimHostFactory from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger from infection_monkey.telemetry.scan_telem import ScanTelem @@ -15,9 +15,15 @@ logger = logging.getLogger() class Propagator: - def __init__(self, telemetry_messenger: ITelemetryMessenger, ip_scanner: IPScanner): + def __init__( + self, + telemetry_messenger: ITelemetryMessenger, + ip_scanner: IPScanner, + victim_host_factory: VictimHostFactory, + ): self._telemetry_messenger = telemetry_messenger self._ip_scanner = ip_scanner + self._victim_host_factory = victim_host_factory self._hosts_to_exploit = None def propagate(self, propagation_config: Dict, stop: Event): @@ -52,7 +58,7 @@ class Propagator: logger.info("Finished network scan") def _process_scan_results(self, ip: str, scan_results: IPScanResults): - victim_host = VictimHost(ip) + victim_host = self._victim_host_factory.build_victim_host(ip) Propagator._process_ping_scan_results(victim_host, scan_results.ping_scan_data) Propagator._process_tcp_scan_results(victim_host, scan_results.port_scan_data) diff --git a/monkey/infection_monkey/model/__init__.py b/monkey/infection_monkey/model/__init__.py index 7c39075be..caf9b6251 100644 --- a/monkey/infection_monkey/model/__init__.py +++ b/monkey/infection_monkey/model/__init__.py @@ -1,4 +1,5 @@ from infection_monkey.model.host import VictimHost +from infection_monkey.model.victim_host_factory import VictimHostFactory MONKEY_ARG = "m0nk3y" DROPPER_ARG = "dr0pp3r" diff --git a/monkey/infection_monkey/model/victim_host_factory.py b/monkey/infection_monkey/model/victim_host_factory.py new file mode 100644 index 000000000..e3ac8d5a7 --- /dev/null +++ b/monkey/infection_monkey/model/victim_host_factory.py @@ -0,0 +1,28 @@ +from infection_monkey.model import VictimHost + + +class VictimHostFactory: + def __init__(self): + pass + + def build_victim_host(self, ip: str): + victim_host = VictimHost(ip) + + # TODO: Reimplement the below logic from the old monkey.py + """ + if self._monkey_tunnel: + self._monkey_tunnel.set_tunnel_for_host(machine) + if self._default_server: + if self._network.on_island(self._default_server): + machine.set_default_server( + get_interface_to_target(machine.ip_addr) + + (":" + self._default_server_port if self._default_server_port else "") + ) + else: + machine.set_default_server(self._default_server) + logger.debug( + f"Default server for machine: {machine} set to {machine.default_server}" + ) + """ + + return victim_host diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py b/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py index 1610e752b..0584ca1cd 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py @@ -2,7 +2,7 @@ from infection_monkey.master import AutomatedMaster def test_terminate_without_start(): - m = AutomatedMaster(None, None, None) + m = AutomatedMaster(None, None, None, None) # Test that call to terminate does not raise exception m.terminate() 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 d8f65b54e..941f17a6c 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_propagator.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_propagator.py @@ -2,6 +2,7 @@ from threading import Event from infection_monkey.i_puppet import FingerprintData, PingScanData, PortScanData, PortStatus from infection_monkey.master import IPScanResults, Propagator +from infection_monkey.model import VictimHostFactory empty_fingerprint_data = FingerprintData(None, None, {}) @@ -87,7 +88,7 @@ class MockIPScanner: def test_scan_result_processing(telemetry_messenger_spy): - p = Propagator(telemetry_messenger_spy, MockIPScanner()) + p = Propagator(telemetry_messenger_spy, MockIPScanner(), VictimHostFactory()) p.propagate( { "targets": {"subnet_scan_list": ["10.0.0.1", "10.0.0.2", "10.0.0.3"]},