p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

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

This commit is contained in:
Shreya Malviya 2022-09-07 18:57:21 +05:30 committed by Mike Salvatore
parent 789d6b8441
commit 18659b654f
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
monkey/infection_monkey/network/relay/utils.py
View File

@ -1,6 +1,6 @@
import logging import logging
import socket import socket
from typing import Sequence from typing import Optional, Sequence
import requests import requests
@ -12,7 +12,9 @@ from infection_monkey.utils.threading import create_daemon_thread
logger = logging.getLogger(__name__) 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)}") logger.debug(f"Trying to wake up with servers: {', '.join(servers)}")
server_iterator = (s for s in servers) server_iterator = (s for s in servers)
@ -26,6 +28,8 @@ def find_server(self, servers: Sequence[str]):
timeout=MEDIUM_REQUEST_TIMEOUT, timeout=MEDIUM_REQUEST_TIMEOUT,
) )
server_found = server
break break
except requests.exceptions.ConnectionError as err: except requests.exceptions.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}")
@ -44,6 +48,8 @@ def find_server(self, servers: Sequence[str]):
) )
t.start() t.start()
return server_found
def _send_relay_control_message(server: str): def _send_relay_control_message(server: str):
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket: