From b2f0b54d37ce8bf69491e32bce711fb200127ca1 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 5 Sep 2022 13:07:38 +0200 Subject: [PATCH] Agent: Add list of server when propagating --- monkey/infection_monkey/master/propagator.py | 16 ++++++++++++++-- .../infection_monkey/master/test_propagator.py | 10 +++++++--- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/master/propagator.py b/monkey/infection_monkey/master/propagator.py index 03a5d5ec7..74a0a41ca 100644 --- a/monkey/infection_monkey/master/propagator.py +++ b/monkey/infection_monkey/master/propagator.py @@ -47,7 +47,11 @@ class Propagator: self._hosts_to_exploit: Queue = Queue() def propagate( - self, propagation_config: PropagationConfiguration, current_depth: int, stop: Event + self, + propagation_config: PropagationConfiguration, + current_depth: int, + servers: Sequence[str], + stop: Event, ): logger.info("Attempting to propagate") @@ -66,7 +70,13 @@ class Propagator: exploit_thread = create_daemon_thread( target=self._exploit_hosts, name="PropagatorExploitThread", - args=(propagation_config.exploitation, current_depth, network_scan_completed, stop), + args=( + propagation_config.exploitation, + current_depth, + servers, + network_scan_completed, + stop, + ), ) scan_thread.start() @@ -167,6 +177,7 @@ class Propagator: self, exploitation_config: ExploitationConfiguration, current_depth: int, + servers: Sequence[str], network_scan_completed: Event, stop: Event, ): @@ -176,6 +187,7 @@ class Propagator: exploitation_config, self._hosts_to_exploit, current_depth, + servers, self._process_exploit_attempts, network_scan_completed, 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 53136e755..8f1b51274 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_propagator.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_propagator.py @@ -110,6 +110,8 @@ os_windows = "windows" os_linux = "linux" +SERVERS = ["127.0.0.1:5000", "10.10.10.10:5007"] + @pytest.fixture def mock_ip_scanner(): @@ -134,6 +136,7 @@ class StubExploiter: exploiters_to_run, hosts_to_exploit, current_depth, + servers, results_callback, scan_completed, stop, @@ -171,7 +174,7 @@ def test_scan_result_processing( subnets=["10.0.0.1", "10.0.0.2", "10.0.0.3"], ) propagation_config = get_propagation_config(default_agent_configuration, targets) - p.propagate(propagation_config, 1, Event()) + p.propagate(propagation_config, 1, SERVERS, Event()) assert len(telemetry_messenger_spy.telemetries) == 3 @@ -204,6 +207,7 @@ class MockExploiter: exploiters_to_run, hosts_to_exploit, current_depth, + servers, results_callback, scan_completed, stop, @@ -269,7 +273,7 @@ def test_exploiter_result_processing( subnets=["10.0.0.1", "10.0.0.2", "10.0.0.3"], ) propagation_config = get_propagation_config(default_agent_configuration, targets) - p.propagate(propagation_config, 1, Event()) + p.propagate(propagation_config, 1, SERVERS, Event()) exploit_telems = [t for t in telemetry_messenger_spy.telemetries if isinstance(t, ExploitTelem)] assert len(exploit_telems) == 4 @@ -310,7 +314,7 @@ def test_scan_target_generation( subnets=["10.0.0.0/29", "172.10.20.30"], ) propagation_config = get_propagation_config(default_agent_configuration, targets) - p.propagate(propagation_config, 1, Event()) + p.propagate(propagation_config, 1, SERVERS, Event()) expected_ip_scan_list = [ "10.0.0.0",