From 18659b654fdea601574573a68fb85829cab42dcd Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 7 Sep 2022 18:57:21 +0530 Subject: [PATCH] Agent: Return server to connect to or None from `find_server` --- monkey/infection_monkey/network/relay/utils.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 56462dc33..66f0b1d1e 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -1,6 +1,6 @@ import logging import socket -from typing import Sequence +from typing import Optional, Sequence import requests @@ -12,7 +12,9 @@ from infection_monkey.utils.threading import create_daemon_thread logger = logging.getLogger(__name__) -def find_server(self, servers: Sequence[str]): +def find_server(self, servers: Sequence[str]) -> Optional[str]: + server_found = None + logger.debug(f"Trying to wake up with servers: {', '.join(servers)}") server_iterator = (s for s in servers) @@ -26,6 +28,8 @@ def find_server(self, servers: Sequence[str]): timeout=MEDIUM_REQUEST_TIMEOUT, ) + server_found = server + break except requests.exceptions.ConnectionError as err: logger.error(f"Unable to connect to server/relay {server}: {err}") @@ -44,6 +48,8 @@ def find_server(self, servers: Sequence[str]): ) t.start() + return server_found + def _send_relay_control_message(server: str): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket: