p15693087
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests 1 Projects Releases Wiki Activity

Agent: Wrap relay stop logic in a method

This commit is contained in:
Mike Salvatore 2022-09-12 14:53:19 -04:00
parent 27cb2009fc
commit 77c97062eb
1 changed files with 11 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
monkey/infection_monkey/monkey.py
View File

@ -393,15 +393,7 @@ class InfectionMonkey:
self._master.cleanup() self._master.cleanup()
reset_signal_handlers() reset_signal_handlers()
self._stop_relay()
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)
if firewall.is_enabled(): if firewall.is_enabled():
firewall.remove_firewall_rule() firewall.remove_firewall_rule()
@ -425,6 +417,16 @@ class InfectionMonkey:
logger.info("Monkey is shutting down") 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): def _close_tunnel(self):
logger.info(f"Quitting tunnel {self._cmd_island_ip}") logger.info(f"Quitting tunnel {self._cmd_island_ip}")
notify_disconnect(self._cmd_island_ip, self._cmd_island_port) notify_disconnect(self._cmd_island_ip, self._cmd_island_port)