From cadc23d8a5ad58294515d00138a67cec336180d6 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 18 Mar 2022 13:02:30 -0400 Subject: [PATCH] Agent: Only start/stop tunnel if the agent is able to propagate Starting and stopping the tunnel is slow, and only necessary if the agent plans to propagate. If depth < 1, propagation will not occur, so there's no point in having a tunnel open. If a `-d` parameter is not supplied to the agent, the tunnel will be started. --- monkey/infection_monkey/monkey.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 5a6238dfc..23d423473 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -154,7 +154,7 @@ class InfectionMonkey: firewall.add_firewall_rule() self._monkey_inbound_tunnel = ControlClient.create_control_tunnel() - if self._monkey_inbound_tunnel: + if self._monkey_inbound_tunnel and self._propagation_enabled(): self._monkey_inbound_tunnel.start() StateTelem(is_done=False, version=get_version()).send() @@ -265,7 +265,7 @@ class InfectionMonkey: reset_signal_handlers() - if self._monkey_inbound_tunnel: + if self._monkey_inbound_tunnel and self._propagation_enabled(): self._monkey_inbound_tunnel.stop() self._monkey_inbound_tunnel.join() @@ -291,6 +291,12 @@ class InfectionMonkey: logger.info("Monkey is shutting down") + def _propagation_enabled(self) -> bool: + # If self._current_depth is None, assume that propagation is desired. + # The Master will ignore this value if it is None and pull the actual + # maximum depth from the server + return self._current_depth is None or self._current_depth > 0 + @staticmethod def _close_tunnel(): tunnel_address = (