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
from ipaddress import IPv4Address
from logging import getLogger
from .relay_user_handler import RelayUserHandler
from .tcp_pipe_spawner import TCPPipeSpawner
RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST = b"infection-monkey-relay-control-message: -"
logger = getLogger(__name__)
class RelayConnectionHandler:
"""Handles new relay connections."""
@ -29,4 +32,7 @@ class RelayConnectionHandler:
self._relay_user_handler.disconnect_user(addr)
else:
self._relay_user_handler.add_relay_user(addr)
self._pipe_spawner.spawn_pipe(sock)
try:
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:
read_list, _, except_list = select.select(sockets, [], sockets, self.timeout)
if except_list:
raise Exception("select() failed on sockets {except_list}")
raise OSError("select() failed on sockets {except_list}")
if not read_list:
raise TimeoutError("pipe did not receive data for {self.timeout} seconds")
@ -46,17 +46,17 @@ class SocketsPipe(Thread):
def run(self):
try:
self._pipe()
except Exception as err:
except OSError as err:
logger.debug(err)
try:
self.source.close()
except Exception as err:
except OSError as err:
logger.debug(f"Error while closing source socket: {err}")
try:
self.dest.close()
except Exception as err:
except OSError as err:
logger.debug(f"Error while closing destination socket: {err}")
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
: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)
try:
dest.connect((self._target_addr, self._target_port))
except socket.error as err:
dest.connect((str(self._target_addr), self._target_port))
except OSError as err:
source.close()
dest.close()
raise err