Agent: Accept list of servers in AutomatedMaster

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

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