From 787af6ae1b83a435586b1afc5195b070bd0e0467 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 08:57:24 -0400 Subject: [PATCH] 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. --- monkey/infection_monkey/monkey.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 9bb9ef62c..b7eed8b0a 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -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