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

Agent: Use iterator in ControlClient.find_server

This commit is contained in:
Ilija Lazoroski 2022-09-07 14:49:32 +02:00 committed by Mike Salvatore
parent 02a919123b
commit 178b296f75
1 changed files with 11 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
monkey/infection_monkey/control.py
View File

@ -70,8 +70,9 @@ class ControlClient:
def find_server(self, servers: Sequence[str]): def find_server(self, servers: Sequence[str]):
logger.debug(f"Trying to wake up with servers: {', '.join(servers)}") logger.debug(f"Trying to wake up with servers: {', '.join(servers)}")
while servers: server_iterator = (s for s in servers)
server = servers[0]
for server in server_iterator:
try: try:
debug_message = f"Trying to connect to server: {server}" debug_message = f"Trying to connect to server: {server}"
@ -82,37 +83,26 @@ class ControlClient:
timeout=MEDIUM_REQUEST_TIMEOUT, timeout=MEDIUM_REQUEST_TIMEOUT,
) )
# We remove the server that has been succesfull break
servers.remove(server)
# TODO: Check how we are going to set the server address that the ControlCLient # TODO: Check how we are going to set the server address that the ControlCLient
# is going to use # is going to use
# self.server_address = server # self.server_address = server
# If we have any other server we send them RELAY_CONTROL_MESSAGE
if servers:
for ss in servers:
t = create_daemon_thread(
target=ControlClient._send_relay_control_message,
name="SendControlRelayMessageThread",
args=(ss,),
)
t.start()
return True
except ConnectionError as err: except ConnectionError as err:
logger.error(f"Unable to connect to server/relay {server}: {err}") logger.error(f"Unable to connect to server/relay {server}: {err}")
servers.remove(server)
except TimeoutError as err: except TimeoutError as err:
logger.error(f"Timed out while connecting to server/relay {server}: {err}") logger.error(f"Timed out while connecting to server/relay {server}: {err}")
servers.remove(server)
except Exception as err: except Exception as err:
logger.error( logger.error(
f"Exception encountered when trying to connect to server/relay {server}: {err}" f"Exception encountered when trying to connect to server/relay {server}: {err}"
) )
servers.remove(server)
return False for server in server_iterator:
t = create_daemon_thread(
target=ControlClient._send_relay_control_message,
name="SendControlRelayMessageThread",
args=(server,),
)
t.start()
@staticmethod @staticmethod
def _send_relay_control_message(server: str): def _send_relay_control_message(server: str):