Agent: Provide address as string to connect()

This commit is contained in:
Kekoa Kaaikala 2022-09-09 17:08:16 +00:00
parent 46152c0cd2
commit 756bfe488b
3 changed files with 14 additions and 8 deletions

View File

@ -1,11 +1,14 @@
import socket import socket
from ipaddress import IPv4Address from ipaddress import IPv4Address
from logging import getLogger
from .relay_user_handler import RelayUserHandler from .relay_user_handler import RelayUserHandler
from .tcp_pipe_spawner import TCPPipeSpawner from .tcp_pipe_spawner import TCPPipeSpawner
RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST = b"infection-monkey-relay-control-message: -" RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST = b"infection-monkey-relay-control-message: -"
logger = getLogger(__name__)
class RelayConnectionHandler: class RelayConnectionHandler:
"""Handles new relay connections.""" """Handles new relay connections."""
@ -29,4 +32,7 @@ class RelayConnectionHandler:
self._relay_user_handler.disconnect_user(addr) self._relay_user_handler.disconnect_user(addr)
else: else:
self._relay_user_handler.add_relay_user(addr) self._relay_user_handler.add_relay_user(addr)
try:
self._pipe_spawner.spawn_pipe(sock) self._pipe_spawner.spawn_pipe(sock)
except OSError as err:
logger.debug(f"Failed to spawn pipe: {err}")

View File

@ -32,7 +32,7 @@ class SocketsPipe(Thread):
while True: while True:
read_list, _, except_list = select.select(sockets, [], sockets, self.timeout) read_list, _, except_list = select.select(sockets, [], sockets, self.timeout)
if except_list: if except_list:
raise Exception("select() failed on sockets {except_list}") raise OSError("select() failed on sockets {except_list}")
if not read_list: if not read_list:
raise TimeoutError("pipe did not receive data for {self.timeout} seconds") raise TimeoutError("pipe did not receive data for {self.timeout} seconds")
@ -46,17 +46,17 @@ class SocketsPipe(Thread):
def run(self): def run(self):
try: try:
self._pipe() self._pipe()
except Exception as err: except OSError as err:
logger.debug(err) logger.debug(err)
try: try:
self.source.close() self.source.close()
except Exception as err: except OSError as err:
logger.debug(f"Error while closing source socket: {err}") logger.debug(f"Error while closing source socket: {err}")
try: try:
self.dest.close() self.dest.close()
except Exception as err: except OSError as err:
logger.debug(f"Error while closing destination socket: {err}") logger.debug(f"Error while closing destination socket: {err}")
self._pipe_closed(self) self._pipe_closed(self)

View File

@ -22,12 +22,12 @@ class TCPPipeSpawner:
Attempt to create a pipe on between the configured client and the provided socket Attempt to create a pipe on between the configured client and the provided socket
:param source: A socket to the connecting client. :param source: A socket to the connecting client.
:raises socket.error: If a socket to the configured client could not be created. :raises OSError: If a socket to the configured client could not be created.
""" """
dest = socket.socket(socket.AF_INET, socket.SOCK_STREAM) dest = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try: try:
dest.connect((self._target_addr, self._target_port)) dest.connect((str(self._target_addr), self._target_port))
except socket.error as err: except OSError as err:
source.close() source.close()
dest.close() dest.close()
raise err raise err