Agent: Use SocketAddress in TCPRelay

monkey/infection_monkey/monkey.py

@@ -21,6 +21,7 @@ from common.network.network_utils import (
     get_my_ip_addresses,
     get_network_interfaces,
 )
+from common.types import SocketAddress
 from common.utils.argparse_types import positive_int
 from common.utils.attack_utils import ScanStatus, UsageEnum
 from common.version import get_version
@@ -230,10 +231,10 @@ class InfectionMonkey:
         config = self._control_channel.get_config()
 
         relay_port = get_free_tcp_port()
+        island_address = SocketAddress(IPv4Address(self._cmd_island_ip), self._cmd_island_port)
         self._relay = TCPRelay(
             relay_port,
-            IPv4Address(self._cmd_island_ip),
-            self._cmd_island_port,
+            island_address,
             client_disconnect_timeout=config.keep_tunnel_open_time,
         )
 

monkey/infection_monkey/network/relay/tcp_relay.py

@@ -3,6 +3,7 @@ from logging import getLogger
 from threading import Lock, Thread
 from time import sleep
 
+from common.types import SocketAddress
 from infection_monkey.network.relay import (
     RelayConnectionHandler,
     RelayUserHandler,
@@ -22,15 +23,14 @@ class TCPRelay(Thread, InterruptableThreadMixin):
     def __init__(
         self,
         relay_port: int,
-        dest_addr: IPv4Address,
-        dest_port: int,
+        dest_address: SocketAddress,
         client_disconnect_timeout: float,
     ):
         self._user_handler = RelayUserHandler(
             new_client_timeout=client_disconnect_timeout,
             client_disconnect_timeout=client_disconnect_timeout,
         )
-        self._pipe_spawner = TCPPipeSpawner(dest_addr, dest_port)
+        self._pipe_spawner = TCPPipeSpawner(dest_address.ip, dest_address.port)
         relay_filter = RelayConnectionHandler(self._pipe_spawner, self._user_handler)
         self._connection_handler = TCPConnectionHandler(
             bind_host="",