Merge pull request #2354 from guardicore/2323-use-socketaddress-in-running_on_island

2323 use socketaddress in running on island

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 (
-    address_to_ip_port,
-    get_my_ip_addresses,
-    get_network_interfaces,
-)
+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
@@ -118,17 +114,17 @@ class InfectionMonkey:
 
         self._agent_event_serializer_registry = self._setup_agent_event_serializers()
 
-        server, self._island_api_client = self._connect_to_island_api()
-        self._cmd_island_ip = server.ip
-        self._cmd_island_port = server.port
-
-        self._island_address = SocketAddress(self._cmd_island_ip, self._cmd_island_port)
+        self._island_address, self._island_api_client = self._connect_to_island_api()
+        self._cmd_island_ip = self._island_address.ip
+        self._cmd_island_port = self._island_address.port
 
         self._control_client = ControlClient(
-            server_address=str(server), island_api_client=self._island_api_client
+            server_address=self._island_address, 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)
+        self._control_channel = ControlChannel(
+            str(self._island_address), get_agent_id(), self._island_api_client
+        )
+        self._register_agent()
 
         # TODO Refactor the telemetry messengers to accept control client
         # and remove control_client_object
@@ -178,14 +174,14 @@ class InfectionMonkey:
 
         return server, island_api_client
 
-    def _register_agent(self, server: SocketAddress):
+    def _register_agent(self):
         agent_registration_data = AgentRegistrationData(
             id=get_agent_id(),
             machine_hardware_id=get_machine_id(),
             start_time=agent_process.get_start_time(),
             # parent_id=parent,
             parent_id=None,  # None for now, until we change GUID to UUID
-            cc_server=server,
+            cc_server=self._island_address,
             network_interfaces=get_network_interfaces(),
         )
         self._island_api_client.register_agent(agent_registration_data)
@@ -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)
-        return server_ip in {str(interface.ip) for interface in local_network_interfaces}
+        server_ip = self._control_client.server_address.ip
+        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],
-                get_interface_to_target(self.control_client.server_address.split(":")[0]),
+                server_ip,
+                get_interface_to_target(server_ip),
                 filename,
             )
         )

monkey/infection_monkey/post_breach/pba.py

@@ -1,6 +1,6 @@
 import logging
 import subprocess
-from typing import Dict, Iterable, List, Tuple
+from typing import Dict, Iterable, List, Optional, Tuple
 
 from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT
 from common.utils.attack_utils import ScanStatus
@@ -24,7 +24,7 @@ class PBA:
         name="unknown",
         linux_cmd="",
         windows_cmd="",
-        timeout: int = LONG_REQUEST_TIMEOUT,
+        timeout: Optional[float] = LONG_REQUEST_TIMEOUT,
     ):
         """
         :param name: Name of post breach action.