Merge pull request #2352 from guardicore/2323-SocketAddress-in-TCPRelay

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
@@ -120,6 +121,11 @@ class InfectionMonkey:
         # TODO: `address_to_port()` should return the port as an integer.
         self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server)
         self._cmd_island_port = int(self._cmd_island_port)
+
+        self._island_address = SocketAddress(
+            IPv4Address(self._cmd_island_ip), self._cmd_island_port
+        )
+
         self._control_client = ControlClient(
             server_address=server, island_api_client=self._island_api_client
         )
@@ -232,8 +238,7 @@ class InfectionMonkey:
         relay_port = get_free_tcp_port()
         self._relay = TCPRelay(
             relay_port,
-            IPv4Address(self._cmd_island_ip),
+            self._island_address,
-            self._cmd_island_port,
             client_disconnect_timeout=config.keep_tunnel_open_time,
         )
 
@@ -487,7 +492,7 @@ class InfectionMonkey:
 
     def _close_tunnel(self):
         logger.info(f"Quitting tunnel {self._cmd_island_ip}")
-        notify_disconnect(self._cmd_island_ip, self._cmd_island_port)
+        notify_disconnect(self._island_address)
 
     def _send_log(self):
         monkey_log_path = get_agent_log_path()

monkey/infection_monkey/network/relay/tcp_pipe_spawner.py

@@ -1,9 +1,10 @@
 import socket
-from ipaddress import IPv4Address
 from logging import getLogger
 from threading import Lock
 from typing import Set
 
+from common.types import SocketAddress
+
 from .consts import SOCKET_TIMEOUT
 from .sockets_pipe import SocketsPipe
 
@@ -15,9 +16,9 @@ class TCPPipeSpawner:
     Creates bi-directional pipes between the configured client and other clients.
     """
 
-    def __init__(self, target_addr: IPv4Address, target_port: int):
+    def __init__(self, target_addr: SocketAddress):
-        self._target_addr = target_addr
+        self._target_ip = target_addr.ip
-        self._target_port = target_port
+        self._target_port = target_addr.port
         self._pipes: Set[SocketsPipe] = set()
         self._lock = Lock()
 
@@ -31,7 +32,7 @@ class TCPPipeSpawner:
         dest = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         dest.settimeout(SOCKET_TIMEOUT)
         try:
-            dest.connect((str(self._target_addr), self._target_port))
+            dest.connect((str(self._target_ip), self._target_port))
         except OSError as err:
             source.close()
             dest.close()

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_address: SocketAddress,
-        dest_port: int,
         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)
         relay_filter = RelayConnectionHandler(self._pipe_spawner, self._user_handler)
         self._connection_handler = TCPConnectionHandler(
             bind_host="",

monkey/infection_monkey/network/relay/utils.py

@@ -6,6 +6,7 @@ from typing import Dict, Iterable, Iterator, Optional
 
 from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT
 from common.network.network_utils import address_to_ip_port
+from common.types import SocketAddress
 from infection_monkey.island_api_client import (
     AbstractIslandAPIClientFactory,
     IIslandAPIClient,
@@ -90,22 +91,22 @@ def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]):
 
 def _send_remove_from_waitlist_control_message_to_relay(server: str):
     ip, port = address_to_ip_port(server)
-    notify_disconnect(IPv4Address(ip), int(port))
+    server_address = SocketAddress(IPv4Address(ip), int(port))
+    notify_disconnect(server_address)
 
 
-def notify_disconnect(server_ip: IPv4Address, server_port: int):
+def notify_disconnect(server_address: SocketAddress):
     """
-    Tell upstream relay that we no longer need the relay.
+    Tell upstream relay that we no longer need the relay
 
-    :param server_ip: The IP address of the server to notify.
+    :param server_address: The address of the server to notify
-    :param server_port: The port of the server to notify.
     """
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket:
         d_socket.settimeout(LONG_REQUEST_TIMEOUT)
 
         try:
-            d_socket.connect((str(server_ip), server_port))
+            d_socket.connect((str(server_address.ip), server_address.port))
             d_socket.sendall(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST)
-            logger.info(f"Control message was sent to the server/relay {server_ip}:{server_port}")
+            logger.info(f"Control message was sent to the server/relay {server_address}")
         except OSError as err:
-            logger.error(f"Error connecting to socket {server_ip}:{server_port}: {err}")
+            logger.error(f"Error connecting to socket {server_address}: {err}")

vulture_allowlist.py

@@ -8,7 +8,6 @@ from common.agent_configuration.agent_sub_configurations import (
     ScanTargetConfiguration,
 )
 from common.credentials import Credentials, LMHash, NTHash
-from common.types import SocketAddress
 from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFactory
 from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue
 from monkey_island.cc.models import Report
@@ -326,6 +325,3 @@ SCANNED
 EXPLOITED
 CC
 CC_TUNNEL
-
-# TODO: Remove after #2323
-SocketAddress