From c33189725dae70dc3a1ec4a41e512476b729a41b Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 26 Sep 2022 15:58:41 +0000 Subject: [PATCH] Agent: Update ControlClient to use SocketAddress --- monkey/infection_monkey/control.py | 11 ++--------- monkey/infection_monkey/monkey.py | 12 ++++-------- monkey/infection_monkey/network/tools.py | 2 +- .../post_breach/custom_pba/custom_pba.py | 5 +++-- 4 files changed, 10 insertions(+), 20 deletions(-) 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 e336e6db9..d7384d679 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 @@ -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) - 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, ) )