From 77c97062eb8e00d0922f33648d29d5cf0cdbb488 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 12 Sep 2022 14:53:19 -0400 Subject: [PATCH] Agent: Wrap relay stop logic in a method --- monkey/infection_monkey/monkey.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index b4d9897ce..9effa91b0 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -393,15 +393,7 @@ class InfectionMonkey: self._master.cleanup() reset_signal_handlers() - - if self._relay and self._relay.is_alive(): - self._relay.stop() - - while self._relay.is_alive() and not self._control_channel.should_agent_stop(): - self._relay.join(timeout=5) - - if self._control_channel.should_agent_stop(): - self._relay.join(timeout=60) + self._stop_relay() if firewall.is_enabled(): firewall.remove_firewall_rule() @@ -425,6 +417,16 @@ class InfectionMonkey: logger.info("Monkey is shutting down") + def _stop_relay(self): + if self._relay and self._relay.is_alive(): + self._relay.stop() + + while self._relay.is_alive() and not self._control_channel.should_agent_stop(): + self._relay.join(timeout=5) + + if self._control_channel.should_agent_stop(): + self._relay.join(timeout=60) + def _close_tunnel(self): logger.info(f"Quitting tunnel {self._cmd_island_ip}") notify_disconnect(self._cmd_island_ip, self._cmd_island_port)