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.
This commit is contained in:
Mike Salvatore 2021-12-13 16:53:07 -05:00
parent 8423a064bb
commit 93d0bb6cd2
6 changed files with 44 additions and 6 deletions

View File

@ -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)

View File

@ -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)

View File

@ -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"

View File

@ -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

View File

@ -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()

View File

@ -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"]},