From 756bfe488b383c2067fa759018e556843321d00a Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Fri, 9 Sep 2022 17:08:16 +0000
Subject: [PATCH] Agent: Provide address as string to connect()

---
 .../network/relay/relay_connection_handler.py             | 8 +++++++-
 monkey/infection_monkey/network/relay/sockets_pipe.py     | 8 ++++----
 monkey/infection_monkey/network/relay/tcp_pipe_spawner.py | 6 +++---
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/monkey/infection_monkey/network/relay/relay_connection_handler.py b/monkey/infection_monkey/network/relay/relay_connection_handler.py
index 4b4475e52..9d081cf65 100644
--- a/monkey/infection_monkey/network/relay/relay_connection_handler.py
+++ b/monkey/infection_monkey/network/relay/relay_connection_handler.py
@@ -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}")
diff --git a/monkey/infection_monkey/network/relay/sockets_pipe.py b/monkey/infection_monkey/network/relay/sockets_pipe.py
index e31dfaaf4..9ec400d0d 100644
--- a/monkey/infection_monkey/network/relay/sockets_pipe.py
+++ b/monkey/infection_monkey/network/relay/sockets_pipe.py
@@ -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)
diff --git a/monkey/infection_monkey/network/relay/tcp_pipe_spawner.py b/monkey/infection_monkey/network/relay/tcp_pipe_spawner.py
index 5ffebaaea..8e702d374 100644
--- a/monkey/infection_monkey/network/relay/tcp_pipe_spawner.py
+++ b/monkey/infection_monkey/network/relay/tcp_pipe_spawner.py
@@ -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