From 524b97078de4c7111f819a8677900afaa740b31b Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Mon, 7 Mar 2022 08:46:10 -0500
Subject: [PATCH] Agent: Pass current depth to AutomatedMaster

---
 monkey/infection_monkey/master/automated_master.py    |  8 ++++++--
 monkey/infection_monkey/monkey.py                     | 11 ++---------
 .../infection_monkey/master/test_automated_master.py  |  6 +++---
 3 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py
index d90c8e902..51627d728 100644
--- a/monkey/infection_monkey/master/automated_master.py
+++ b/monkey/infection_monkey/master/automated_master.py
@@ -1,7 +1,7 @@
 import logging
 import threading
 import time
-from typing import Any, Callable, Dict, Iterable, List, Tuple
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
 
 from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
 from infection_monkey.i_master import IMaster
@@ -30,12 +30,14 @@ logger = logging.getLogger()
 class AutomatedMaster(IMaster):
     def __init__(
         self,
+        current_depth: Optional[int],
         puppet: IPuppet,
         telemetry_messenger: ITelemetryMessenger,
         victim_host_factory: VictimHostFactory,
         control_channel: IControlChannel,
         local_network_interfaces: List[NetworkInterface],
     ):
+        self._current_depth = current_depth
         self._puppet = puppet
         self._telemetry_messenger = telemetry_messenger
         self._control_channel = control_channel
@@ -162,7 +164,9 @@ class AutomatedMaster(IMaster):
         # still running.
         credential_collector_thread.join()
 
-        current_depth = config["depth"]
+        current_depth = self._current_depth if self._current_depth is not None else config["depth"]
+        logger.info(f"Current depth is {current_depth}")
+
         if self._can_propagate() and current_depth > 0:
             self._propagator.propagate(config["propagation"], current_depth, self._stop)
 
diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index 155289a31..db32c9703 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/monkey/infection_monkey/monkey.py
@@ -68,6 +68,7 @@ class InfectionMonkey:
         # TODO used in propogation phase
         self._monkey_inbound_tunnel = None
         self.telemetry_messenger = LegacyTelemetryMessengerAdapter()
+        self._current_depth = self._opts.depth
 
     @staticmethod
     def _get_arguments(args):
@@ -93,7 +94,6 @@ class InfectionMonkey:
 
         logger.info("Monkey is starting...")
 
-        self._set_propagation_depth(self._opts)
         self._add_default_server_to_config(self._opts.server)
         self._connect_to_island()
 
@@ -111,14 +111,6 @@ class InfectionMonkey:
         self._setup()
         self._master.start()
 
-    @staticmethod
-    def _set_propagation_depth(options):
-        if options.depth is not None:
-            WormConfiguration._depth_from_commandline = True
-            WormConfiguration.depth = options.depth
-            logger.debug("Setting propagation depth from command line")
-        logger.debug(f"Set propagation depth to {WormConfiguration.depth}")
-
     @staticmethod
     def _add_default_server_to_config(default_server: str):
         if default_server:
@@ -179,6 +171,7 @@ class InfectionMonkey:
         )
 
         self._master = AutomatedMaster(
+            self._current_depth,
             puppet,
             telemetry_messenger,
             victim_host_factory,
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 c7023e525..4bb7b4294 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, MagicMock(), [])
+    m = AutomatedMaster(None, None, None, None, 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, cc, [])
+    m = AutomatedMaster(None, None, None, None, cc, [])
     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, cc, [])
+    m = AutomatedMaster(None, None, None, None, cc, [])
     m.start()
 
     assert cc.should_agent_stop.call_count == CHECK_FOR_STOP_AGENT_COUNT