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

Agent: Let relay run indefinitely if user hasn't forced a stop

This commit is contained in:
Mike Salvatore 2022-09-12 14:46:42 -04:00
parent c2c2993ff7
commit 27cb2009fc
1 changed files with 12 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
monkey/infection_monkey/monkey.py
View File

@ -110,6 +110,7 @@ class InfectionMonkey:
# TODO Refactor the telemetry messengers to accept control client # TODO Refactor the telemetry messengers to accept control client
# and remove control_client_object # and remove control_client_object
ControlClient.control_client_object = self._control_client ControlClient.control_client_object = self._control_client
self._control_channel = None
self._telemetry_messenger = LegacyTelemetryMessengerAdapter() self._telemetry_messenger = LegacyTelemetryMessengerAdapter()
self._current_depth = self._opts.depth self._current_depth = self._opts.depth
self._master = None self._master = None
@ -177,10 +178,10 @@ class InfectionMonkey:
if firewall.is_enabled(): if firewall.is_enabled():
firewall.add_firewall_rule() firewall.add_firewall_rule()
control_channel = ControlChannel(self._control_client.server_address, GUID) self._control_channel = ControlChannel(self._control_client.server_address, GUID)
control_channel.register_agent(self._opts.parent) self._control_channel.register_agent(self._opts.parent)
config = control_channel.get_config() config = self._control_channel.get_config()
relay_port = get_free_tcp_port() relay_port = get_free_tcp_port()
self._relay = TCPRelay( self._relay = TCPRelay(
@ -204,9 +205,8 @@ class InfectionMonkey:
local_network_interfaces = InfectionMonkey._get_local_network_interfaces() local_network_interfaces = InfectionMonkey._get_local_network_interfaces()
# TODO control_channel and control_client have same responsibilities, merge them # TODO control_channel and control_client have same responsibilities, merge them
control_channel = ControlChannel(self._control_client.server_address, GUID)
propagation_credentials_repository = AggregatingPropagationCredentialsRepository( propagation_credentials_repository = AggregatingPropagationCredentialsRepository(
control_channel self._control_channel
) )
event_queue = PyPubSubAgentEventQueue(Publisher()) event_queue = PyPubSubAgentEventQueue(Publisher())
@ -226,7 +226,7 @@ class InfectionMonkey:
puppet, puppet,
telemetry_messenger, telemetry_messenger,
victim_host_factory, victim_host_factory,
control_channel, self._control_channel,
local_network_interfaces, local_network_interfaces,
propagation_credentials_repository, propagation_credentials_repository,
) )
@ -396,7 +396,12 @@ class InfectionMonkey:
if self._relay and self._relay.is_alive(): if self._relay and self._relay.is_alive():
self._relay.stop() self._relay.stop()
self._relay.join(timeout=60)
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()