diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 45b329ad1..6959e6bf1 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -2,7 +2,7 @@ import logging import threading import time from ipaddress import IPv4Interface -from typing import Any, Callable, Collection, List, Optional +from typing import Any, Callable, Collection, List, Optional, Sequence from common.agent_configuration import CustomPBAConfiguration, PluginConfiguration from common.utils import Timer @@ -35,6 +35,7 @@ class AutomatedMaster(IMaster): def __init__( self, current_depth: Optional[int], + servers: Sequence[str], puppet: IPuppet, telemetry_messenger: ITelemetryMessenger, victim_host_factory: VictimHostFactory, @@ -43,6 +44,7 @@ class AutomatedMaster(IMaster): credentials_store: IPropagationCredentialsRepository, ): self._current_depth = current_depth + self._servers = servers self._puppet = puppet self._telemetry_messenger = telemetry_messenger self._control_channel = control_channel @@ -175,7 +177,7 @@ class AutomatedMaster(IMaster): logger.info(f"Current depth is {current_depth}") if maximum_depth_reached(config.propagation.maximum_depth, current_depth): - self._propagator.propagate(config.propagation, current_depth, self._stop) + self._propagator.propagate(config.propagation, current_depth, self._servers, self._stop) else: logger.info("Skipping propagation: maximum depth reached") 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 9029ce480..1ebcc886f 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 @@ -14,7 +14,7 @@ INTERVAL = 0.001 def test_terminate_without_start(): - m = AutomatedMaster(None, None, None, None, MagicMock(), [], MagicMock()) + m = AutomatedMaster(None, [], None, None, None, MagicMock(), [], MagicMock()) # Test that call to terminate does not raise exception m.terminate() @@ -34,7 +34,7 @@ def test_stop_if_cant_get_config_from_island(monkeypatch): monkeypatch.setattr( "infection_monkey.master.automated_master.CHECK_FOR_TERMINATE_INTERVAL_SEC", INTERVAL ) - m = AutomatedMaster(None, None, None, None, cc, [], MagicMock()) + m = AutomatedMaster(None, [], None, None, None, cc, [], MagicMock()) m.start() assert cc.get_config.call_count == CHECK_FOR_CONFIG_COUNT @@ -73,7 +73,7 @@ def test_stop_if_cant_get_stop_signal_from_island(monkeypatch, sleep_and_return_ "infection_monkey.master.automated_master.CHECK_FOR_TERMINATE_INTERVAL_SEC", INTERVAL ) - m = AutomatedMaster(None, None, None, None, cc, [], MagicMock()) + m = AutomatedMaster(None, [], None, None, None, cc, [], MagicMock()) m.start() assert cc.should_agent_stop.call_count == CHECK_FOR_STOP_AGENT_COUNT