diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 4b0361608..41b3511d9 100644 --- a/monkey/infection_monkey/control.py +++ b/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)) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index d8e4962d9..acd10ce95 100644 --- a/monkey/infection_monkey/monkey.py +++ b/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() diff --git a/monkey/infection_monkey/network/tools.py b/monkey/infection_monkey/network/tools.py index c612a7e48..2a309956c 100644 --- a/monkey/infection_monkey/network/tools.py +++ b/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.' diff --git a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py index 34fb73147..64276be6b 100644 --- a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py +++ b/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, ) ) diff --git a/monkey/infection_monkey/post_breach/pba.py b/monkey/infection_monkey/post_breach/pba.py index c381f6ad9..fd89a47d8 100644 --- a/monkey/infection_monkey/post_breach/pba.py +++ b/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.