Agent: Return server to connect to or None from `find_server`

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: