Agent: Update ControlClient to use SocketAddress

monkey/infection_monkey/control.py

@@ -8,6 +8,7 @@ from urllib3 import disable_warnings
 
 from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT
 from common.network.network_utils import get_my_ip_addresses_legacy
+from common.types import SocketAddress
 from infection_monkey.config import GUID
 from infection_monkey.island_api_client import IIslandAPIClient
 from infection_monkey.network.info import get_host_subnets
@@ -24,7 +25,7 @@ class ControlClient:
     # https://github.com/guardicore/monkey/blob/133f7f5da131b481561141171827d1f9943f6aec/monkey/infection_monkey/telemetry/base_telem.py
     control_client_object = None
 
-    def __init__(self, server_address: str, island_api_client: IIslandAPIClient):
+    def __init__(self, server_address: SocketAddress, island_api_client: IIslandAPIClient):
         self.server_address = server_address
         self._island_api_client = island_api_client
 
@@ -55,12 +56,6 @@ class ControlClient:
         )
 
     def send_telemetry(self, telem_category, json_data: str):
-        if not self.server_address:
-            logger.error(
-                "Trying to send %s telemetry before current server is established, aborting."
-                % telem_category
-            )
-            return
         try:
             telemetry = {"monkey_guid": GUID, "telem_category": telem_category, "data": json_data}
             requests.post(  # noqa: DUO123
@@ -74,8 +69,6 @@ class ControlClient:
             logger.warning(f"Error connecting to control server {self.server_address}: {exc}")
 
     def send_log(self, log):
-        if not self.server_address:
-            return
         try:
             telemetry = {"monkey_guid": GUID, "log": json.dumps(log)}
             self._island_api_client.send_log(json.dumps(telemetry))

monkey/infection_monkey/monkey.py

@@ -16,11 +16,7 @@ from common.agent_event_serializers import (
 from common.agent_events import CredentialsStolenEvent
 from common.agent_registration_data import AgentRegistrationData
 from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue
-from common.network.network_utils import (
+from common.network.network_utils import get_my_ip_addresses, get_network_interfaces
-    address_to_ip_port,
-    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
@@ -125,7 +121,7 @@ class InfectionMonkey:
         self._island_address = SocketAddress(self._cmd_island_ip, self._cmd_island_port)
 
         self._control_client = ControlClient(
-            server_address=str(server), island_api_client=self._island_api_client
+            server_address=server, island_api_client=self._island_api_client
         )
         self._control_channel = ControlChannel(str(server), get_agent_id(), self._island_api_client)
         self._register_agent(self._island_address)
@@ -444,8 +440,8 @@ class InfectionMonkey:
         return VictimHostFactory(self._cmd_island_ip, self._cmd_island_port, on_island)
 
     def _running_on_island(self, local_network_interfaces: List[IPv4Interface]) -> bool:
-        server_ip, _ = address_to_ip_port(self._control_client.server_address)
+        server_ip = self._control_client.server_address.ip
-        return server_ip in {str(interface.ip) for interface in local_network_interfaces}
+        return server_ip in {interface.ip for interface in local_network_interfaces}
 
     def _is_another_monkey_running(self):
         return not self._singleton.try_lock()

monkey/infection_monkey/network/tools.py

@@ -51,7 +51,7 @@ def tcp_port_to_service(port):
     return "tcp-" + str(port)
 
 
-def get_interface_to_target(dst):
+def get_interface_to_target(dst: str) -> str:
     """
     :param dst: destination IP address string without port. E.G. '192.168.1.1.'
     :return: IP address string of an interface that can connect to the target. E.G. '192.168.1.4.'

monkey/infection_monkey/post_breach/custom_pba/custom_pba.py

@@ -83,11 +83,12 @@ class CustomPBA(PBA):
         if not status:
             status = ScanStatus.USED
 
+        server_ip = str(self.control_client.server_address.ip)
         self.telemetry_messenger.send_telemetry(
             T1105Telem(
                 status,
-                self.control_client.server_address.split(":")[0],
+                server_ip,
-                get_interface_to_target(self.control_client.server_address.split(":")[0]),
+                get_interface_to_target(server_ip),
                 filename,
             )
         )