From 93d0bb6cd21ae324209d239f47a06a4b0e188d91 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 13 Dec 2021 16:53:07 -0500 Subject: [PATCH] Agent: Add a placeholder VictimHostFactory The AutomatedMaster will need access to the monkey's tunnel, IP addresses, and default server in order to properly configure the victim host. The VictimHostFactory can abstract these dependencies away and handle these details on behalf of the AutomatedMaster. --- .../master/automated_master.py | 4 ++- monkey/infection_monkey/master/propagator.py | 12 ++++++-- monkey/infection_monkey/model/__init__.py | 1 + .../model/victim_host_factory.py | 28 +++++++++++++++++++ .../master/test_automated_master.py | 2 +- .../master/test_propagator.py | 3 +- 6 files changed, 44 insertions(+), 6 deletions(-) create mode 100644 monkey/infection_monkey/model/victim_host_factory.py 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"]},