diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 51f675d11..780e07b6f 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/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] diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 6e3d9b0f0..8b8d8e600 100644 --- a/monkey/infection_monkey/monkey.py +++ b/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() diff --git a/monkey/infection_monkey/utils/propagation.py b/monkey/infection_monkey/utils/propagation.py new file mode 100644 index 000000000..568def5d5 --- /dev/null +++ b/monkey/infection_monkey/utils/propagation.py @@ -0,0 +1,2 @@ +def should_propagate(config: dict, depth: int) -> bool: + return config["config"]["depth"] < depth