Agent: Fix send relay disconnect to unneeded relays

Since `find_server()` is parallelized, the iterator was completely
exhausted when `send_remove_from_waitlist_control_message_to_relays()`
was called, making it effectively a NOOP.
This commit is contained in:
Mike Salvatore 2022-09-19 08:57:24 -04:00 committed by Ilija Lazoroski
parent 9456a30bd9
commit 787af6ae1b
1 changed files with 4 additions and 4 deletions

View File

@ -138,8 +138,7 @@ class InfectionMonkey:
def _get_server(self):
logger.debug(f"Trying to wake up with servers: {', '.join(self._opts.servers)}")
servers_iterator = (s for s in self._opts.servers)
server = find_server(servers_iterator)
server = find_server(self._opts.servers)
if server:
logger.info(f"Successfully connected to the island via {server}")
else:
@ -147,10 +146,11 @@ class InfectionMonkey:
f"Failed to connect to the island via any known servers: {self._opts.servers}"
)
# Note: Since we pass the address for each of our interfaces to the exploited
# NOTE: Since we pass the address for each of our interfaces to the exploited
# machines, is it possible for a machine to unintentionally unregister itself from the
# relay if it is able to connect to the relay over multiple interfaces?
send_remove_from_waitlist_control_message_to_relays(servers_iterator)
servers_to_close = (s for s in self._opts.servers if s != server)
send_remove_from_waitlist_control_message_to_relays(servers_to_close)
return server