Agent: Create and use should_propagate method

This method is used to determine if the inbound tunnel should be open and if the master should attempt exploiting other machines

monkey/infection_monkey/master/automated_master.py

@@ -14,6 +14,7 @@ from infection_monkey.telemetry.credentials_telem import CredentialsTelem
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
 from infection_monkey.utils.threading import create_daemon_thread, interruptible_iter
+from utils.propagation import should_propagate
 
 from . import Exploiter, IPScanner, Propagator
 from .option_parsing import custom_pba_is_enabled
@@ -172,7 +173,7 @@ class AutomatedMaster(IMaster):
         current_depth = self._current_depth if self._current_depth is not None else 0
         logger.info(f"Current depth is {current_depth}")
 
-        if self._can_propagate() and current_depth < config["depth"]:
+        if self._can_propagate():
             self._propagator.propagate(config["propagation"], current_depth, self._stop)
 
         payload_thread = create_daemon_thread(
@@ -201,7 +202,7 @@ class AutomatedMaster(IMaster):
             self._telemetry_messenger.send_telemetry(PostBreachTelem(pba_data))
 
     def _can_propagate(self) -> bool:
-        return True
+        return should_propagate(self._control_channel.get_config(), self._current_depth)
 
     def _run_payload(self, payload: Tuple[str, Dict]):
         name = payload[0]

monkey/infection_monkey/monkey.py

@@ -79,6 +79,7 @@ from infection_monkey.utils.monkey_dir import (
 )
 from infection_monkey.utils.monkey_log_path import get_agent_log_path
 from infection_monkey.utils.signal_handler import register_signal_handlers, reset_signal_handlers
+from utils.propagation import should_propagate
 
 logger = logging.getLogger(__name__)
 logging.getLogger("urllib3").setLevel(logging.INFO)
@@ -167,7 +168,10 @@ class InfectionMonkey:
             firewall.add_firewall_rule()
 
         self._monkey_inbound_tunnel = self._control_client.create_control_tunnel()
-        if self._monkey_inbound_tunnel:
+        config = ControlChannel(
+            self._control_client.server_address, GUID, self._control_client.proxies
+        ).get_config()
+        if self._monkey_inbound_tunnel and should_propagate(config, self._current_depth):
             self._monkey_inbound_tunnel.start()
 
         StateTelem(is_done=False, version=get_version()).send()

monkey/infection_monkey/utils/propagation.py

@@ -0,0 +1,2 @@
+def should_propagate(config: dict, depth: int) -> bool:
+    return config["config"]["depth"] < depth