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.
This commit is contained in:
Mike Salvatore 2022-03-18 13:02:30 -04:00
parent fe7c7d5d9c
commit cadc23d8a5
1 changed files with 8 additions and 2 deletions

View File

@ -154,7 +154,7 @@ class InfectionMonkey:
firewall.add_firewall_rule() firewall.add_firewall_rule()
self._monkey_inbound_tunnel = ControlClient.create_control_tunnel() 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() self._monkey_inbound_tunnel.start()
StateTelem(is_done=False, version=get_version()).send() StateTelem(is_done=False, version=get_version()).send()
@ -265,7 +265,7 @@ class InfectionMonkey:
reset_signal_handlers() 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.stop()
self._monkey_inbound_tunnel.join() self._monkey_inbound_tunnel.join()
@ -291,6 +291,12 @@ class InfectionMonkey:
logger.info("Monkey is shutting down") 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 @staticmethod
def _close_tunnel(): def _close_tunnel():
tunnel_address = ( tunnel_address = (