From 049eb1b174449b367755193fdde2fbbdf29cdcf9 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Mon, 13 Jun 2022 13:14:48 +0300 Subject: [PATCH 01/17] Agent: Add control client to the agent initialization --- monkey/infection_monkey/control.py | 63 ++++++++++++++---------------- monkey/infection_monkey/monkey.py | 49 +++++++++++------------ 2 files changed, 52 insertions(+), 60 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 77b52fe3f..338a8183e 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -3,6 +3,7 @@ import logging import platform from pprint import pformat from socket import gethostname +from typing import Mapping, Optional import requests from requests.exceptions import ConnectionError @@ -23,11 +24,12 @@ logger = logging.getLogger(__name__) PBA_FILE_DOWNLOAD = "https://%s/api/pba/download/%s" -class ControlClient(object): - proxies = {} +class ControlClient: + def __init__(self, server_address: str, proxies: Optional[Mapping[str, str]] = None): + self.proxies = {} if not proxies else proxies + self.server_address = server_address - @staticmethod - def wakeup(parent=None): + def wakeup(self, parent=None): if parent: logger.debug("parent: %s" % (parent,)) @@ -45,20 +47,19 @@ class ControlClient(object): "launch_time": agent_process.get_start_time(), } - if ControlClient.proxies: - monkey["tunnel"] = ControlClient.proxies.get("https") + if self.proxies: + monkey["tunnel"] = self.proxies.get("https") requests.post( # noqa: DUO123 f"https://{WormConfiguration.current_server}/api/agent", data=json.dumps(monkey), headers={"content-type": "application/json"}, verify=False, - proxies=ControlClient.proxies, + proxies=self.proxies, timeout=MEDIUM_REQUEST_TIMEOUT, ) - @staticmethod - def find_server(default_tunnel=None): + def find_server(self, default_tunnel=None): logger.debug( "Trying to wake up with Monkey Island servers list: %r" % WormConfiguration.command_servers @@ -73,13 +74,13 @@ class ControlClient(object): current_server = server debug_message = "Trying to connect to server: %s" % server - if ControlClient.proxies: - debug_message += " through proxies: %s" % ControlClient.proxies + if self.proxies: + debug_message += " through proxies: %s" % self.proxies logger.debug(debug_message) requests.get( # noqa: DUO123 f"https://{server}/api?action=is-up", verify=False, - proxies=ControlClient.proxies, + proxies=self.proxies, timeout=MEDIUM_REQUEST_TIMEOUT, ) WormConfiguration.current_server = current_server @@ -92,20 +93,19 @@ class ControlClient(object): if current_server: return True else: - if ControlClient.proxies: + if self.proxies: return False else: logger.info("Starting tunnel lookup...") proxy_find = tunnel.find_tunnel(default=default_tunnel) if proxy_find: - ControlClient.set_proxies(proxy_find) - return ControlClient.find_server() + self.set_proxies(proxy_find) + return self.find_server() else: logger.info("No tunnel found") return False - @staticmethod - def set_proxies(proxy_find): + def set_proxies(self, proxy_find): """ Note: The proxy schema changes between different versions of requests and urllib3, which causes the machine to not open a tunnel back. @@ -120,12 +120,11 @@ class ControlClient(object): proxy_address, proxy_port = proxy_find logger.info("Found tunnel at %s:%s" % (proxy_address, proxy_port)) if is_windows_os(): - ControlClient.proxies["https"] = f"http://{proxy_address}:{proxy_port}" + self.proxies["https"] = f"http://{proxy_address}:{proxy_port}" else: - ControlClient.proxies["https"] = f"{proxy_address}:{proxy_port}" + self.proxies["https"] = f"{proxy_address}:{proxy_port}" - @staticmethod - def send_telemetry(telem_category, json_data: str): + def send_telemetry(self, telem_category, json_data: str): if not WormConfiguration.current_server: logger.error( "Trying to send %s telemetry before current server is established, aborting." @@ -139,7 +138,7 @@ class ControlClient(object): data=json.dumps(telemetry), headers={"content-type": "application/json"}, verify=False, - proxies=ControlClient.proxies, + proxies=self.proxies, timeout=MEDIUM_REQUEST_TIMEOUT, ) except Exception as exc: @@ -147,8 +146,7 @@ class ControlClient(object): "Error connecting to control server %s: %s", WormConfiguration.current_server, exc ) - @staticmethod - def send_log(log): + def send_log(self, log): if not WormConfiguration.current_server: return try: @@ -158,7 +156,7 @@ class ControlClient(object): data=json.dumps(telemetry), headers={"content-type": "application/json"}, verify=False, - proxies=ControlClient.proxies, + proxies=self.proxies, timeout=MEDIUM_REQUEST_TIMEOUT, ) except Exception as exc: @@ -166,15 +164,14 @@ class ControlClient(object): "Error connecting to control server %s: %s", WormConfiguration.current_server, exc ) - @staticmethod - def load_control_config(): + def load_control_config(self): if not WormConfiguration.current_server: return try: reply = requests.get( # noqa: DUO123 f"https://{WormConfiguration.current_server}/api/agent/", verify=False, - proxies=ControlClient.proxies, + proxies=self.proxies, timeout=MEDIUM_REQUEST_TIMEOUT, ) @@ -200,12 +197,11 @@ class ControlClient(object): ) raise Exception("Couldn't load from from server's configuration, aborting. %s" % exc) - @staticmethod - def create_control_tunnel(): + def create_control_tunnel(self): if not WormConfiguration.current_server: return None - my_proxy = ControlClient.proxies.get("https", "").replace("https://", "") + my_proxy = self.proxies.get("https", "").replace("https://", "") if my_proxy: proxy_class = TcpProxy try: @@ -224,13 +220,12 @@ class ControlClient(object): target_port=target_port, ) - @staticmethod - def get_pba_file(filename): + def get_pba_file(self, filename): try: return requests.get( # noqa: DUO123 PBA_FILE_DOWNLOAD % (WormConfiguration.current_server, filename), verify=False, - proxies=ControlClient.proxies, + proxies=self.proxies, timeout=LONG_REQUEST_TIMEOUT, ) except requests.exceptions.RequestException: diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 4d2807194..13930ecda 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -89,7 +89,7 @@ class InfectionMonkey: self._singleton = SystemSingleton() self._opts = self._get_arguments(args) self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.server) - self._default_server = self._opts.server + self.cc_client = ControlClient(self._opts.server) self._monkey_inbound_tunnel = None self._telemetry_messenger = LegacyTelemetryMessengerAdapter() self._current_depth = self._opts.depth @@ -120,7 +120,6 @@ class InfectionMonkey: logger.info("Agent is starting...") logger.info(f"Agent GUID: {GUID}") - self._add_default_server_to_config(self._opts.server) self._connect_to_island() # TODO: Reevaluate who is responsible to send this information @@ -137,27 +136,20 @@ class InfectionMonkey: self._setup() self._master.start() - @staticmethod - def _add_default_server_to_config(default_server: str): - if default_server: - logger.debug("Added default server: %s" % default_server) - WormConfiguration.command_servers.insert(0, default_server) - def _connect_to_island(self): # Sets island's IP and port for monkey to communicate to if self._current_server_is_set(): - self._default_server = WormConfiguration.current_server - logger.debug("Default server set to: %s" % self._default_server) + logger.debug("Default server set to: %s" % self.cc_client.server_address) else: raise Exception( "Monkey couldn't find server with {} default tunnel.".format(self._opts.tunnel) ) - ControlClient.wakeup(parent=self._opts.parent) - ControlClient.load_control_config() + self.cc_client.wakeup(parent=self._opts.parent) + self.cc_client.load_control_config() def _current_server_is_set(self) -> bool: - if ControlClient.find_server(default_tunnel=self._opts.tunnel): + if self.cc_client.find_server(default_tunnel=self._opts.tunnel): return True return False @@ -170,7 +162,7 @@ class InfectionMonkey: if firewall.is_enabled(): firewall.add_firewall_rule() - self._monkey_inbound_tunnel = ControlClient.create_control_tunnel() + self._monkey_inbound_tunnel = self.cc_client.create_control_tunnel() if self._monkey_inbound_tunnel and self._propagation_enabled(): self._monkey_inbound_tunnel.start() @@ -184,7 +176,9 @@ class InfectionMonkey: def _build_master(self): local_network_interfaces = InfectionMonkey._get_local_network_interfaces() - control_channel = ControlChannel(self._default_server, GUID) + # TODO control_channel and control_client have same responsibilities, merge them + control_channel = ControlChannel(self.cc_client.server_address, GUID) + control_client = self.cc_client credentials_store = AggregatingCredentialsStore(control_channel) puppet = self._build_puppet(credentials_store) @@ -206,6 +200,7 @@ class InfectionMonkey: control_channel, local_network_interfaces, credentials_store, + control_client, ) @staticmethod @@ -237,7 +232,7 @@ class InfectionMonkey: puppet.load_plugin("ssh", SSHFingerprinter(), PluginType.FINGERPRINTER) agent_repository = CachingAgentRepository( - f"https://{self._default_server}", ControlClient.proxies + f"https://{self.cc_client.server_address}", self.cc_client.proxies ) exploit_wrapper = ExploiterWrapper(self._telemetry_messenger, agent_repository) @@ -320,7 +315,9 @@ class InfectionMonkey: PluginType.POST_BREACH_ACTION, ) puppet.load_plugin( - "CustomPBA", CustomPBA(self._telemetry_messenger), PluginType.POST_BREACH_ACTION + "CustomPBA", + CustomPBA(self._telemetry_messenger, self.cc_client), + PluginType.POST_BREACH_ACTION, ) puppet.load_plugin("ransomware", RansomwarePayload(), PluginType.PAYLOAD) @@ -338,7 +335,7 @@ class InfectionMonkey: ) def _running_on_island(self, local_network_interfaces: List[NetworkInterface]) -> bool: - server_ip, _ = address_to_ip_port(self._default_server) + server_ip, _ = address_to_ip_port(self.cc_client.server_address) return server_ip in {interface.address for interface in local_network_interfaces} def _is_another_monkey_running(self): @@ -363,13 +360,13 @@ class InfectionMonkey: deleted = InfectionMonkey._self_delete() - InfectionMonkey._send_log() + self._send_log() StateTelem( is_done=True, version=get_version() ).send() # Signal the server (before closing the tunnel) - InfectionMonkey._close_tunnel() + self._close_tunnel() self._singleton.unlock() except Exception as e: logger.error(f"An error occurred while cleaning up the monkey agent: {e}") @@ -384,15 +381,15 @@ class InfectionMonkey: # maximum depth from the server return self._current_depth is None or self._current_depth > 0 - @staticmethod - def _close_tunnel(): - tunnel_address = ControlClient.proxies.get("https", "").replace("http://", "").split(":")[0] + def _close_tunnel(self): + tunnel_address = ( + self.cc_client.proxies.get("https", "").replace("http://", "").split(":")[0] + ) if tunnel_address: logger.info("Quitting tunnel %s", tunnel_address) tunnel.quit_tunnel(tunnel_address) - @staticmethod - def _send_log(): + def _send_log(self): monkey_log_path = get_agent_log_path() if monkey_log_path.is_file(): with open(monkey_log_path, "r") as f: @@ -400,7 +397,7 @@ class InfectionMonkey: else: log = "" - ControlClient.send_log(log) + self.cc_client.send_log(log) @staticmethod def _self_delete() -> bool: From a099f21f61fd104ddaadc6194beb055a7688e5cb Mon Sep 17 00:00:00 2001 From: vakarisz Date: Mon, 13 Jun 2022 13:17:09 +0300 Subject: [PATCH 02/17] Agent: Initialize CustomPBA with a ControlClient object This is done to refactor ControlClient from a global --- monkey/infection_monkey/post_breach/custom_pba/custom_pba.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 ce7dd64f2..f86df9e28 100644 --- a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py +++ b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py @@ -25,11 +25,12 @@ class CustomPBA(PBA): Defines user's configured post breach action. """ - def __init__(self, telemetry_messenger: ITelemetryMessenger): + def __init__(self, telemetry_messenger: ITelemetryMessenger, cc_client: ControlClient): super(CustomPBA, self).__init__( telemetry_messenger, POST_BREACH_FILE_EXECUTION, timeout=None ) self.filename = "" + self.cc_client = cc_client def run(self, options: Dict) -> Iterable[PostBreachData]: self._set_options(options) @@ -75,7 +76,7 @@ class CustomPBA(PBA): :return: True if successful, false otherwise """ - pba_file_contents = ControlClient.get_pba_file(filename) + pba_file_contents = self.cc_client.get_pba_file(filename) status = None if not pba_file_contents or not pba_file_contents.content: From 799ff3d6fd4a88847378d8adb4bdb4bda324db4d Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 13 Jun 2022 12:54:17 +0200 Subject: [PATCH 03/17] UT: Fix control client and custom pba tests --- .../actions/test_users_custom_pba.py | 30 ++++++++++++------- .../infection_monkey/test_control.py | 4 +-- 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py index 598f3412b..75da530b0 100644 --- a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py +++ b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py @@ -2,6 +2,7 @@ from unittest.mock import MagicMock import pytest +from infection_monkey.control import ControlClient from infection_monkey.post_breach.custom_pba.custom_pba import CustomPBA MONKEY_DIR_PATH = "/dir/to/monkey/" @@ -48,8 +49,15 @@ def fake_custom_pba_linux_options(): } -def test_command_linux_custom_file_and_cmd(fake_custom_pba_linux_options, set_os_linux): - pba = CustomPBA(MagicMock()) +@pytest.fixture +def control_client(): + return ControlClient(CUSTOM_SERVER) + + +def test_command_linux_custom_file_and_cmd( + fake_custom_pba_linux_options, set_os_linux, control_client +): + pba = CustomPBA(MagicMock(), control_client) pba._set_options(fake_custom_pba_linux_options) expected_command = f"cd {MONKEY_DIR_PATH} ; {CUSTOM_LINUX_CMD}" assert pba.command == expected_command @@ -68,9 +76,11 @@ def fake_custom_pba_windows_options(): } -def test_command_windows_custom_file_and_cmd(fake_custom_pba_windows_options, set_os_windows): +def test_command_windows_custom_file_and_cmd( + fake_custom_pba_windows_options, set_os_windows, control_client +): - pba = CustomPBA(MagicMock()) + pba = CustomPBA(MagicMock(), control_client) pba._set_options(fake_custom_pba_windows_options) expected_command = f"cd {MONKEY_DIR_PATH} & {CUSTOM_WINDOWS_CMD}" assert pba.command == expected_command @@ -90,8 +100,8 @@ def fake_options_files_only(): @pytest.mark.parametrize("os", [set_os_linux, set_os_windows]) -def test_files_only(fake_options_files_only, os): - pba = CustomPBA(MagicMock()) +def test_files_only(fake_options_files_only, os, control_client): + pba = CustomPBA(MagicMock(), control_client) pba._set_options(fake_options_files_only) assert pba.command == "" @@ -108,15 +118,15 @@ def fake_options_commands_only(): } -def test_commands_only(fake_options_commands_only, set_os_linux): - pba = CustomPBA(MagicMock()) +def test_commands_only(fake_options_commands_only, set_os_linux, control_client): + pba = CustomPBA(MagicMock(), control_client) pba._set_options(fake_options_commands_only) assert pba.command == CUSTOM_LINUX_CMD assert pba.filename == "" -def test_commands_only_windows(fake_options_commands_only, set_os_windows): - pba = CustomPBA(MagicMock()) +def test_commands_only_windows(fake_options_commands_only, set_os_windows, control_client): + pba = CustomPBA(MagicMock(), control_client) pba._set_options(fake_options_commands_only) assert pba.command == CUSTOM_WINDOWS_CMD assert pba.filename == "" diff --git a/monkey/tests/unit_tests/infection_monkey/test_control.py b/monkey/tests/unit_tests/infection_monkey/test_control.py index 7c2f42d19..6310f68c0 100644 --- a/monkey/tests/unit_tests/infection_monkey/test_control.py +++ b/monkey/tests/unit_tests/infection_monkey/test_control.py @@ -8,8 +8,8 @@ from monkey.infection_monkey.control import ControlClient [(True, "http://8.8.8.8:45455"), (False, "8.8.8.8:45455")], ) def test_control_set_proxies(monkeypatch, is_windows_os, expected_proxy_string): - monkeypatch.setattr("monkey.infection_monkey.control.is_windows_os", lambda: is_windows_os) - control_client = ControlClient() + monkeypatch.setattr("infection_monkey.cc.control.is_windows_os", lambda: is_windows_os) + control_client = ControlClient("8.8.8.8:5000") control_client.set_proxies(("8.8.8.8", "45455")) From fb1a577823122826325035bd416b85f7adb3acd6 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 13 Jun 2022 12:55:41 +0200 Subject: [PATCH 04/17] Agent: Add control client proxies to tunnel telem initialization --- monkey/infection_monkey/monkey.py | 2 +- monkey/infection_monkey/telemetry/tunnel_telem.py | 7 ++++--- .../infection_monkey/telemetry/test_tunnel_telem.py | 2 +- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 13930ecda..e8be1f05d 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -167,7 +167,7 @@ class InfectionMonkey: self._monkey_inbound_tunnel.start() StateTelem(is_done=False, version=get_version()).send() - TunnelTelem().send() + TunnelTelem(self.cc_client.proxies).send() self._build_master() diff --git a/monkey/infection_monkey/telemetry/tunnel_telem.py b/monkey/infection_monkey/telemetry/tunnel_telem.py index f8d562771..efe917643 100644 --- a/monkey/infection_monkey/telemetry/tunnel_telem.py +++ b/monkey/infection_monkey/telemetry/tunnel_telem.py @@ -1,15 +1,16 @@ +from typing import Mapping + from common.common_consts.telem_categories import TelemCategoryEnum -from infection_monkey.control import ControlClient from infection_monkey.telemetry.base_telem import BaseTelem class TunnelTelem(BaseTelem): - def __init__(self): + def __init__(self, proxy: Mapping[str, str]): """ Default tunnel telemetry constructor """ super(TunnelTelem, self).__init__() - self.proxy = ControlClient.proxies.get("https") + self.proxy = proxy.get("https") telem_category = TelemCategoryEnum.TUNNEL diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/test_tunnel_telem.py b/monkey/tests/unit_tests/infection_monkey/telemetry/test_tunnel_telem.py index eab763790..eb18307ce 100644 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/test_tunnel_telem.py +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/test_tunnel_telem.py @@ -7,7 +7,7 @@ from infection_monkey.telemetry.tunnel_telem import TunnelTelem @pytest.fixture def tunnel_telem_test_instance(): - return TunnelTelem() + return TunnelTelem({}) def test_tunnel_telem_send(tunnel_telem_test_instance, spy_send_telemetry): From 94dbd9a8e2be24f4178f749d8cc80fb6a60480df Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 13 Jun 2022 13:02:17 +0200 Subject: [PATCH 05/17] Agent: Add proxies to the initialization of ControlChannel --- monkey/infection_monkey/master/control_channel.py | 10 ++++++---- monkey/infection_monkey/monkey.py | 8 ++++++-- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 028437faf..0834f28f7 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -1,5 +1,6 @@ import json import logging +from typing import Mapping import requests @@ -14,9 +15,10 @@ logger = logging.getLogger(__name__) class ControlChannel(IControlChannel): - def __init__(self, server: str, agent_id: str): + def __init__(self, server: str, agent_id: str, proxies: Mapping[str, str]): self._agent_id = agent_id self._control_channel_server = server + self._proxies = proxies def should_agent_stop(self) -> bool: if not self._control_channel_server: @@ -30,7 +32,7 @@ class ControlChannel(IControlChannel): response = requests.get( # noqa: DUO123 url, verify=False, - proxies=ControlClient.proxies, + proxies=self._proxies, timeout=SHORT_REQUEST_TIMEOUT, ) response.raise_for_status() @@ -51,7 +53,7 @@ class ControlChannel(IControlChannel): response = requests.get( # noqa: DUO123 f"https://{self._control_channel_server}/api/agent", verify=False, - proxies=ControlClient.proxies, + proxies=self._proxies, timeout=SHORT_REQUEST_TIMEOUT, ) response.raise_for_status() @@ -74,7 +76,7 @@ class ControlChannel(IControlChannel): response = requests.get( # noqa: DUO123 propagation_credentials_url, verify=False, - proxies=ControlClient.proxies, + proxies=self._proxies, timeout=SHORT_REQUEST_TIMEOUT, ) response.raise_for_status() diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index e8be1f05d..eda19632a 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -128,7 +128,9 @@ class InfectionMonkey: run_aws_environment_check(self._telemetry_messenger) - should_stop = ControlChannel(WormConfiguration.current_server, GUID).should_agent_stop() + should_stop = ControlChannel( + WormConfiguration.current_server, GUID, self.cc_client.proxies + ).should_agent_stop() if should_stop: logger.info("The Monkey Island has instructed this agent to stop") return @@ -177,7 +179,9 @@ class InfectionMonkey: local_network_interfaces = InfectionMonkey._get_local_network_interfaces() # TODO control_channel and control_client have same responsibilities, merge them - control_channel = ControlChannel(self.cc_client.server_address, GUID) + control_channel = ControlChannel( + self.cc_client.server_address, GUID, self.cc_client.proxies + ) control_client = self.cc_client credentials_store = AggregatingCredentialsStore(control_channel) From c467dde145faae0da824a93ecfc05908b93c5f5b Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 13 Jun 2022 13:58:12 +0200 Subject: [PATCH 06/17] Agent: Add control_client_object to ControlClient * Workaround global class attribute --- monkey/infection_monkey/control.py | 5 +++++ monkey/infection_monkey/monkey.py | 2 ++ monkey/infection_monkey/telemetry/base_telem.py | 2 +- .../tests/unit_tests/infection_monkey/telemetry/conftest.py | 4 +++- 4 files changed, 11 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 338a8183e..7ec41337d 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -25,6 +25,11 @@ PBA_FILE_DOWNLOAD = "https://%s/api/pba/download/%s" class ControlClient: + # TODO Every telemetry should have its own control client + # for the moment that is a big refactor. + # Ref: infection_monkey.telemetry.base_telem.py + control_client_object = None + def __init__(self, server_address: str, proxies: Optional[Mapping[str, str]] = None): self.proxies = {} if not proxies else proxies self.server_address = server_address diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index eda19632a..e4ee07e9c 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -90,6 +90,8 @@ class InfectionMonkey: self._opts = self._get_arguments(args) self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.server) self.cc_client = ControlClient(self._opts.server) + # TODO Refactor the BaseTelem to have its own control client + ControlClient.control_client_object = self.cc_client self._monkey_inbound_tunnel = None self._telemetry_messenger = LegacyTelemetryMessengerAdapter() self._current_depth = self._opts.depth diff --git a/monkey/infection_monkey/telemetry/base_telem.py b/monkey/infection_monkey/telemetry/base_telem.py index 2f8f68892..18e8d2158 100644 --- a/monkey/infection_monkey/telemetry/base_telem.py +++ b/monkey/infection_monkey/telemetry/base_telem.py @@ -35,7 +35,7 @@ class BaseTelem(ITelem, metaclass=abc.ABCMeta): data = self.get_data() serialized_data = json.dumps(data, cls=self.json_encoder) self._log_telem_sending(serialized_data, log_data) - ControlClient.send_telemetry(self.telem_category, serialized_data) + ControlClient.control_client_object.send_telemetry(self.telem_category, serialized_data) @property def json_encoder(self): diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py b/monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py index cbb1b8074..1bcc2e423 100644 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py @@ -11,5 +11,7 @@ def spy_send_telemetry(monkeypatch): _spy_send_telemetry.telem_category = None _spy_send_telemetry.data = None - monkeypatch.setattr(ControlClient, "send_telemetry", _spy_send_telemetry) + control_client = ControlClient("localhost:5000") + ControlClient.control_client_object = control_client + monkeypatch.setattr(ControlClient.control_client_object, "send_telemetry", _spy_send_telemetry) return _spy_send_telemetry From df116e4fb778d8efd505637bd15140e6d1f4a8d8 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 13 Jun 2022 14:23:36 +0200 Subject: [PATCH 07/17] Agent: Remove current_server and command_servers from WormConfiguration --- monkey/infection_monkey/config.py | 6 +-- monkey/infection_monkey/control.py | 63 +++++++++++------------------- monkey/infection_monkey/monkey.py | 4 +- 3 files changed, 26 insertions(+), 47 deletions(-) diff --git a/monkey/infection_monkey/config.py b/monkey/infection_monkey/config.py index 9bdb35d76..8e783dbf5 100644 --- a/monkey/infection_monkey/config.py +++ b/monkey/infection_monkey/config.py @@ -8,7 +8,7 @@ SENSITIVE_FIELDS = [ "exploit_user_list", "exploit_ssh_keys", ] -LOCAL_CONFIG_VARS = ["name", "id", "current_server", "max_depth"] +LOCAL_CONFIG_VARS = ["name", "id", "max_depth"] HIDDEN_FIELD_REPLACEMENT_CONTENT = "hidden" @@ -62,10 +62,6 @@ class Configuration(object): # depth of propagation depth = 2 max_depth = None - current_server = "" - - # Configuration servers to try to connect to, in this order. - command_servers = [] keep_tunnel_open_time = 30 diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 7ec41337d..93a760631 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -56,7 +56,7 @@ class ControlClient: monkey["tunnel"] = self.proxies.get("https") requests.post( # noqa: DUO123 - f"https://{WormConfiguration.current_server}/api/agent", + f"https://{self.server_address}/api/agent", data=json.dumps(monkey), headers={"content-type": "application/json"}, verify=False, @@ -65,37 +65,26 @@ class ControlClient: ) def find_server(self, default_tunnel=None): - logger.debug( - "Trying to wake up with Monkey Island servers list: %r" - % WormConfiguration.command_servers - ) + logger.debug(f"Trying to wake up with Monkey Island server: {self.server_address}") if default_tunnel: logger.debug("default_tunnel: %s" % (default_tunnel,)) - current_server = "" - - for server in WormConfiguration.command_servers: - try: - current_server = server - - debug_message = "Trying to connect to server: %s" % server - if self.proxies: - debug_message += " through proxies: %s" % self.proxies + try: + debug_message = "Trying to connect to server: %s" % self.server_address + if self.proxies: + debug_message += " through proxies: %s" % self.proxies logger.debug(debug_message) requests.get( # noqa: DUO123 - f"https://{server}/api?action=is-up", + f"https://{self.server_address}/api?action=is-up", verify=False, proxies=self.proxies, timeout=MEDIUM_REQUEST_TIMEOUT, ) - WormConfiguration.current_server = current_server - break + except ConnectionError as exc: + self.server_address = "" + logger.warning("Error connecting to control server %s: %s", self.server_address, exc) - except ConnectionError as exc: - current_server = "" - logger.warning("Error connecting to control server %s: %s", server, exc) - - if current_server: + if self.server_address: return True else: if self.proxies: @@ -130,7 +119,7 @@ class ControlClient: self.proxies["https"] = f"{proxy_address}:{proxy_port}" def send_telemetry(self, telem_category, json_data: str): - if not WormConfiguration.current_server: + if not self.server_address: logger.error( "Trying to send %s telemetry before current server is established, aborting." % telem_category @@ -139,7 +128,7 @@ class ControlClient: try: telemetry = {"monkey_guid": GUID, "telem_category": telem_category, "data": json_data} requests.post( # noqa: DUO123 - "https://%s/api/telemetry" % (WormConfiguration.current_server,), + "https://%s/api/telemetry" % (self.server_address,), data=json.dumps(telemetry), headers={"content-type": "application/json"}, verify=False, @@ -147,17 +136,15 @@ class ControlClient: timeout=MEDIUM_REQUEST_TIMEOUT, ) except Exception as exc: - logger.warning( - "Error connecting to control server %s: %s", WormConfiguration.current_server, exc - ) + logger.warning(f"Error connecting to control server {self.server_address}: {exc}") def send_log(self, log): - if not WormConfiguration.current_server: + if not self.server_address: return try: telemetry = {"monkey_guid": GUID, "log": json.dumps(log)} requests.post( # noqa: DUO123 - "https://%s/api/log" % (WormConfiguration.current_server,), + "https://%s/api/log" % (self.server_address,), data=json.dumps(telemetry), headers={"content-type": "application/json"}, verify=False, @@ -165,25 +152,21 @@ class ControlClient: timeout=MEDIUM_REQUEST_TIMEOUT, ) except Exception as exc: - logger.warning( - "Error connecting to control server %s: %s", WormConfiguration.current_server, exc - ) + logger.warning(f"Error connecting to control server {self.server_address}: {exc}") def load_control_config(self): - if not WormConfiguration.current_server: + if not self.server_address: return try: reply = requests.get( # noqa: DUO123 - f"https://{WormConfiguration.current_server}/api/agent/", + f"https://{self.server_address}/api/agent/", verify=False, proxies=self.proxies, timeout=MEDIUM_REQUEST_TIMEOUT, ) except Exception as exc: - logger.warning( - "Error connecting to control server %s: %s", WormConfiguration.current_server, exc - ) + logger.warning(f"Error connecting to control server {self.server_address}: {exc}") return try: @@ -196,14 +179,14 @@ class ControlClient: # we don't continue with default conf here because it might be dangerous logger.error( "Error parsing JSON reply from control server %s (%s): %s", - WormConfiguration.current_server, + self.server_address, reply._content, exc, ) raise Exception("Couldn't load from from server's configuration, aborting. %s" % exc) def create_control_tunnel(self): - if not WormConfiguration.current_server: + if not self.server_address: return None my_proxy = self.proxies.get("https", "").replace("https://", "") @@ -228,7 +211,7 @@ class ControlClient: def get_pba_file(self, filename): try: return requests.get( # noqa: DUO123 - PBA_FILE_DOWNLOAD % (WormConfiguration.current_server, filename), + PBA_FILE_DOWNLOAD % (self.server_address, filename), verify=False, proxies=self.proxies, timeout=LONG_REQUEST_TIMEOUT, diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index e4ee07e9c..07be6ed44 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -10,8 +10,8 @@ import infection_monkey.tunnel as tunnel from common.network.network_utils import address_to_ip_port from common.utils.attack_utils import ScanStatus, UsageEnum from common.version import get_version -from infection_monkey.config import GUID, WormConfiguration from infection_monkey.control import ControlClient +from infection_monkey.config import GUID from infection_monkey.credential_collectors import ( MimikatzCredentialCollector, SSHCredentialCollector, @@ -131,7 +131,7 @@ class InfectionMonkey: run_aws_environment_check(self._telemetry_messenger) should_stop = ControlChannel( - WormConfiguration.current_server, GUID, self.cc_client.proxies + self.cc_client.server_address, GUID, self.cc_client.proxies ).should_agent_stop() if should_stop: logger.info("The Monkey Island has instructed this agent to stop") From 02a30e695065e9f2369d677721191c08978aced3 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 13 Jun 2022 16:28:09 +0200 Subject: [PATCH 08/17] Agent: Remove `current_server` option from custom_pba --- .../infection_monkey/post_breach/custom_pba/custom_pba.py | 7 ++----- .../post_breach/actions/test_users_custom_pba.py | 8 -------- 2 files changed, 2 insertions(+), 13 deletions(-) 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 f86df9e28..5f0b3a9bf 100644 --- a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py +++ b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py @@ -37,9 +37,6 @@ class CustomPBA(PBA): return super().run(options) def _set_options(self, options: Dict): - # Required for attack telemetry - self.current_server = options["current_server"] - if is_windows_os(): # Add windows commands to PBA's if options["windows_filename"]: @@ -89,8 +86,8 @@ class CustomPBA(PBA): self.telemetry_messenger.send_telemetry( T1105Telem( status, - self.current_server.split(":")[0], - get_interface_to_target(self.current_server.split(":")[0]), + self.cc_client.server_address.split(":")[0], + get_interface_to_target(self.cc_client.server_address.split(":")[0]), filename, ) ) diff --git a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py index 75da530b0..b03c098dd 100644 --- a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py +++ b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py @@ -44,8 +44,6 @@ def fake_custom_pba_linux_options(): "linux_filename": CUSTOM_LINUX_FILENAME, "windows_command": "", "windows_filename": "", - # Current server is used for attack telemetry - "current_server": CUSTOM_SERVER, } @@ -71,8 +69,6 @@ def fake_custom_pba_windows_options(): "linux_filename": "", "windows_command": CUSTOM_WINDOWS_CMD, "windows_filename": CUSTOM_WINDOWS_FILENAME, - # Current server is used for attack telemetry - "current_server": CUSTOM_SERVER, } @@ -94,8 +90,6 @@ def fake_options_files_only(): "linux_filename": CUSTOM_LINUX_FILENAME, "windows_command": "", "windows_filename": CUSTOM_WINDOWS_FILENAME, - # Current server is used for attack telemetry - "current_server": CUSTOM_SERVER, } @@ -113,8 +107,6 @@ def fake_options_commands_only(): "linux_filename": "", "windows_command": CUSTOM_WINDOWS_CMD, "windows_filename": "", - # Current server is used for attack telemetry - "current_server": CUSTOM_SERVER, } From 3c8530cf14827dbc0a8fbaf2c6dfc1b282a26c62 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 13 Jun 2022 17:40:01 +0200 Subject: [PATCH 09/17] Agent: Rename cc_client to control_client --- monkey/infection_monkey/monkey.py | 32 +++++++++---------- .../post_breach/custom_pba/custom_pba.py | 10 +++--- 2 files changed, 21 insertions(+), 21 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 07be6ed44..927793c9e 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -89,9 +89,9 @@ class InfectionMonkey: self._singleton = SystemSingleton() self._opts = self._get_arguments(args) self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.server) - self.cc_client = ControlClient(self._opts.server) + self._control_client = ControlClient(self._opts.server) # TODO Refactor the BaseTelem to have its own control client - ControlClient.control_client_object = self.cc_client + ControlClient.control_client_object = self._control_client self._monkey_inbound_tunnel = None self._telemetry_messenger = LegacyTelemetryMessengerAdapter() self._current_depth = self._opts.depth @@ -131,7 +131,7 @@ class InfectionMonkey: run_aws_environment_check(self._telemetry_messenger) should_stop = ControlChannel( - self.cc_client.server_address, GUID, self.cc_client.proxies + self._control_client.server_address, GUID, self._control_client.proxies ).should_agent_stop() if should_stop: logger.info("The Monkey Island has instructed this agent to stop") @@ -143,17 +143,17 @@ class InfectionMonkey: def _connect_to_island(self): # Sets island's IP and port for monkey to communicate to if self._current_server_is_set(): - logger.debug("Default server set to: %s" % self.cc_client.server_address) + logger.debug("Default server set to: %s" % self._control_client.server_address) else: raise Exception( "Monkey couldn't find server with {} default tunnel.".format(self._opts.tunnel) ) - self.cc_client.wakeup(parent=self._opts.parent) - self.cc_client.load_control_config() + self._control_client.wakeup(parent=self._opts.parent) + self._control_client.load_control_config() def _current_server_is_set(self) -> bool: - if self.cc_client.find_server(default_tunnel=self._opts.tunnel): + if self._control_client.find_server(default_tunnel=self._opts.tunnel): return True return False @@ -166,12 +166,12 @@ class InfectionMonkey: if firewall.is_enabled(): firewall.add_firewall_rule() - self._monkey_inbound_tunnel = self.cc_client.create_control_tunnel() + self._monkey_inbound_tunnel = self._control_client.create_control_tunnel() if self._monkey_inbound_tunnel and self._propagation_enabled(): self._monkey_inbound_tunnel.start() StateTelem(is_done=False, version=get_version()).send() - TunnelTelem(self.cc_client.proxies).send() + TunnelTelem(self._control_client.proxies).send() self._build_master() @@ -182,9 +182,9 @@ class InfectionMonkey: # TODO control_channel and control_client have same responsibilities, merge them control_channel = ControlChannel( - self.cc_client.server_address, GUID, self.cc_client.proxies + self._control_client.server_address, GUID, self._control_client.proxies ) - control_client = self.cc_client + control_client = self._control_client credentials_store = AggregatingCredentialsStore(control_channel) puppet = self._build_puppet(credentials_store) @@ -238,7 +238,7 @@ class InfectionMonkey: puppet.load_plugin("ssh", SSHFingerprinter(), PluginType.FINGERPRINTER) agent_repository = CachingAgentRepository( - f"https://{self.cc_client.server_address}", self.cc_client.proxies + f"https://{self._control_client.server_address}", self._control_client.proxies ) exploit_wrapper = ExploiterWrapper(self._telemetry_messenger, agent_repository) @@ -322,7 +322,7 @@ class InfectionMonkey: ) puppet.load_plugin( "CustomPBA", - CustomPBA(self._telemetry_messenger, self.cc_client), + CustomPBA(self._telemetry_messenger, self._control_client), PluginType.POST_BREACH_ACTION, ) @@ -341,7 +341,7 @@ class InfectionMonkey: ) def _running_on_island(self, local_network_interfaces: List[NetworkInterface]) -> bool: - server_ip, _ = address_to_ip_port(self.cc_client.server_address) + server_ip, _ = address_to_ip_port(self._control_client.server_address) return server_ip in {interface.address for interface in local_network_interfaces} def _is_another_monkey_running(self): @@ -389,7 +389,7 @@ class InfectionMonkey: def _close_tunnel(self): tunnel_address = ( - self.cc_client.proxies.get("https", "").replace("http://", "").split(":")[0] + self._control_client.proxies.get("https", "").replace("http://", "").split(":")[0] ) if tunnel_address: logger.info("Quitting tunnel %s", tunnel_address) @@ -403,7 +403,7 @@ class InfectionMonkey: else: log = "" - self.cc_client.send_log(log) + self._control_client.send_log(log) @staticmethod def _self_delete() -> bool: 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 5f0b3a9bf..f5087a91a 100644 --- a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py +++ b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py @@ -25,12 +25,12 @@ class CustomPBA(PBA): Defines user's configured post breach action. """ - def __init__(self, telemetry_messenger: ITelemetryMessenger, cc_client: ControlClient): + def __init__(self, telemetry_messenger: ITelemetryMessenger, control_client: ControlClient): super(CustomPBA, self).__init__( telemetry_messenger, POST_BREACH_FILE_EXECUTION, timeout=None ) self.filename = "" - self.cc_client = cc_client + self.control_client = control_client def run(self, options: Dict) -> Iterable[PostBreachData]: self._set_options(options) @@ -73,7 +73,7 @@ class CustomPBA(PBA): :return: True if successful, false otherwise """ - pba_file_contents = self.cc_client.get_pba_file(filename) + pba_file_contents = self.control_client.get_pba_file(filename) status = None if not pba_file_contents or not pba_file_contents.content: @@ -86,8 +86,8 @@ class CustomPBA(PBA): self.telemetry_messenger.send_telemetry( T1105Telem( status, - self.cc_client.server_address.split(":")[0], - get_interface_to_target(self.cc_client.server_address.split(":")[0]), + self.control_client.server_address.split(":")[0], + get_interface_to_target(self.control_client.server_address.split(":")[0]), filename, ) ) From 7fe6c170cdbdb23f86bfac22e5d894d74d6b830c Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 13 Jun 2022 17:51:41 +0200 Subject: [PATCH 10/17] UT: Fix control client location --- monkey/tests/unit_tests/infection_monkey/test_control.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/tests/unit_tests/infection_monkey/test_control.py b/monkey/tests/unit_tests/infection_monkey/test_control.py index 6310f68c0..b90087ebf 100644 --- a/monkey/tests/unit_tests/infection_monkey/test_control.py +++ b/monkey/tests/unit_tests/infection_monkey/test_control.py @@ -8,7 +8,7 @@ from monkey.infection_monkey.control import ControlClient [(True, "http://8.8.8.8:45455"), (False, "8.8.8.8:45455")], ) def test_control_set_proxies(monkeypatch, is_windows_os, expected_proxy_string): - monkeypatch.setattr("infection_monkey.cc.control.is_windows_os", lambda: is_windows_os) + monkeypatch.setattr("monkey.infection_monkey.control.is_windows_os", lambda: is_windows_os) control_client = ControlClient("8.8.8.8:5000") control_client.set_proxies(("8.8.8.8", "45455")) From 17a0be2fa04c8c38bb590d0bee807ec2b05ee343 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 13 Jun 2022 17:55:19 +0200 Subject: [PATCH 11/17] Agent: Fix control_client_object TODOs --- monkey/infection_monkey/control.py | 4 ++-- monkey/infection_monkey/monkey.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 93a760631..0f1c749e6 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -25,8 +25,8 @@ PBA_FILE_DOWNLOAD = "https://%s/api/pba/download/%s" class ControlClient: - # TODO Every telemetry should have its own control client - # for the moment that is a big refactor. + # TODO When we have mechanism that support telemetry messenger + # with control clients, then this needs to be removed # Ref: infection_monkey.telemetry.base_telem.py control_client_object = None diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 927793c9e..e3eafddce 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -10,8 +10,8 @@ import infection_monkey.tunnel as tunnel from common.network.network_utils import address_to_ip_port from common.utils.attack_utils import ScanStatus, UsageEnum from common.version import get_version -from infection_monkey.control import ControlClient from infection_monkey.config import GUID +from infection_monkey.control import ControlClient from infection_monkey.credential_collectors import ( MimikatzCredentialCollector, SSHCredentialCollector, @@ -90,7 +90,8 @@ class InfectionMonkey: self._opts = self._get_arguments(args) self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.server) self._control_client = ControlClient(self._opts.server) - # TODO Refactor the BaseTelem to have its own control client + # TODO Refactor the telemetry messengers to accept control client + # and remove control_client_object ControlClient.control_client_object = self._control_client self._monkey_inbound_tunnel = None self._telemetry_messenger = LegacyTelemetryMessengerAdapter() From bbcac3217200a5deee22fd76bfdd4811be5e9372 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 13 Jun 2022 18:30:48 +0200 Subject: [PATCH 12/17] Agent: Remove unused control channel from build_master --- monkey/infection_monkey/master/control_channel.py | 1 - monkey/infection_monkey/monkey.py | 2 -- 2 files changed, 3 deletions(-) diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 0834f28f7..c93af43e0 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -5,7 +5,6 @@ from typing import Mapping import requests from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT -from infection_monkey.control import ControlClient from infection_monkey.custom_types import PropagationCredentials from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index e3eafddce..436b36f67 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -185,7 +185,6 @@ class InfectionMonkey: control_channel = ControlChannel( self._control_client.server_address, GUID, self._control_client.proxies ) - control_client = self._control_client credentials_store = AggregatingCredentialsStore(control_channel) puppet = self._build_puppet(credentials_store) @@ -207,7 +206,6 @@ class InfectionMonkey: control_channel, local_network_interfaces, credentials_store, - control_client, ) @staticmethod From 444b34d548816badb1b67759d7a8571784e245bb Mon Sep 17 00:00:00 2001 From: vakarisz Date: Tue, 14 Jun 2022 12:14:27 +0300 Subject: [PATCH 13/17] UT: Change send telemetry spy syntax in conftest.py --- .../tests/unit_tests/infection_monkey/telemetry/conftest.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py b/monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py index 1bcc2e423..ce790f2d9 100644 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py @@ -1,3 +1,5 @@ +from unittest.mock import MagicMock + import pytest from infection_monkey.control import ControlClient @@ -13,5 +15,5 @@ def spy_send_telemetry(monkeypatch): _spy_send_telemetry.data = None control_client = ControlClient("localhost:5000") ControlClient.control_client_object = control_client - monkeypatch.setattr(ControlClient.control_client_object, "send_telemetry", _spy_send_telemetry) + ControlClient.control_client_object.send_telemetry = MagicMock(side_effect=_spy_send_telemetry) return _spy_send_telemetry From ec2d736984990ace8c0b9c9da91c9902a2d3a8a6 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 14 Jun 2022 12:55:33 +0200 Subject: [PATCH 14/17] Agent: Add github permalink to BaseTelem in ControlClient --- monkey/infection_monkey/control.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 0f1c749e6..e30c003b0 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -27,7 +27,7 @@ PBA_FILE_DOWNLOAD = "https://%s/api/pba/download/%s" class ControlClient: # TODO When we have mechanism that support telemetry messenger # with control clients, then this needs to be removed - # Ref: infection_monkey.telemetry.base_telem.py + # https://github.com/guardicore/monkey/blob/133f7f5da131b481561141171827d1f9943f6aec/monkey/infection_monkey/telemetry/base_telem.py control_client_object = None def __init__(self, server_address: str, proxies: Optional[Mapping[str, str]] = None): From f1bc5f47078aa221bf52f4c7c1f196ffb503afc2 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 14 Jun 2022 08:26:38 -0400 Subject: [PATCH 15/17] Agent: Use f-strings in _connect_to_island() --- monkey/infection_monkey/monkey.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 436b36f67..d7a051193 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -144,11 +144,9 @@ class InfectionMonkey: def _connect_to_island(self): # Sets island's IP and port for monkey to communicate to if self._current_server_is_set(): - logger.debug("Default server set to: %s" % self._control_client.server_address) + logger.debug(f"Default server set to: {self._control_client.server_address}") else: - raise Exception( - "Monkey couldn't find server with {} default tunnel.".format(self._opts.tunnel) - ) + raise Exception(f"Monkey couldn't find server with {self._opts.tunnel} default tunnel.") self._control_client.wakeup(parent=self._opts.parent) self._control_client.load_control_config() From e6e6587f46a0a54c77ae8c04560b1ef90dc83942 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Tue, 14 Jun 2022 16:48:44 +0300 Subject: [PATCH 16/17] Agent: Fix bugs in control.py Bugs happened because of incorrect indentation in the recent refactoring attempting to remove worm config dependency --- monkey/infection_monkey/control.py | 37 ++++++++++++++---------------- 1 file changed, 17 insertions(+), 20 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index e30c003b0..985778b7c 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -73,31 +73,28 @@ class ControlClient: debug_message = "Trying to connect to server: %s" % self.server_address if self.proxies: debug_message += " through proxies: %s" % self.proxies - logger.debug(debug_message) - requests.get( # noqa: DUO123 - f"https://{self.server_address}/api?action=is-up", - verify=False, - proxies=self.proxies, - timeout=MEDIUM_REQUEST_TIMEOUT, - ) + logger.debug(debug_message) + requests.get( # noqa: DUO123 + f"https://{self.server_address}/api?action=is-up", + verify=False, + proxies=self.proxies, + timeout=MEDIUM_REQUEST_TIMEOUT, + ) + return True except ConnectionError as exc: - self.server_address = "" logger.warning("Error connecting to control server %s: %s", self.server_address, exc) - if self.server_address: - return True + if self.proxies: + return False else: - if self.proxies: - return False + logger.info("Starting tunnel lookup...") + proxy_find = tunnel.find_tunnel(default=default_tunnel) + if proxy_find: + self.set_proxies(proxy_find) + return self.find_server() else: - logger.info("Starting tunnel lookup...") - proxy_find = tunnel.find_tunnel(default=default_tunnel) - if proxy_find: - self.set_proxies(proxy_find) - return self.find_server() - else: - logger.info("No tunnel found") - return False + logger.info("No tunnel found") + return False def set_proxies(self, proxy_find): """ From 5ff617b811df7defe1f662460a32b1d072579e3a Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 14 Jun 2022 16:49:45 +0200 Subject: [PATCH 17/17] UT: Pass MagicMock instead of instance of ControlClient --- .../actions/test_users_custom_pba.py | 30 +++++++------------ .../infection_monkey/telemetry/conftest.py | 3 +- 2 files changed, 11 insertions(+), 22 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py index b03c098dd..d73fd7cda 100644 --- a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py +++ b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py @@ -2,7 +2,6 @@ from unittest.mock import MagicMock import pytest -from infection_monkey.control import ControlClient from infection_monkey.post_breach.custom_pba.custom_pba import CustomPBA MONKEY_DIR_PATH = "/dir/to/monkey/" @@ -47,15 +46,8 @@ def fake_custom_pba_linux_options(): } -@pytest.fixture -def control_client(): - return ControlClient(CUSTOM_SERVER) - - -def test_command_linux_custom_file_and_cmd( - fake_custom_pba_linux_options, set_os_linux, control_client -): - pba = CustomPBA(MagicMock(), control_client) +def test_command_linux_custom_file_and_cmd(fake_custom_pba_linux_options, set_os_linux): + pba = CustomPBA(MagicMock(), MagicMock()) pba._set_options(fake_custom_pba_linux_options) expected_command = f"cd {MONKEY_DIR_PATH} ; {CUSTOM_LINUX_CMD}" assert pba.command == expected_command @@ -72,11 +64,9 @@ def fake_custom_pba_windows_options(): } -def test_command_windows_custom_file_and_cmd( - fake_custom_pba_windows_options, set_os_windows, control_client -): +def test_command_windows_custom_file_and_cmd(fake_custom_pba_windows_options, set_os_windows): - pba = CustomPBA(MagicMock(), control_client) + pba = CustomPBA(MagicMock(), MagicMock()) pba._set_options(fake_custom_pba_windows_options) expected_command = f"cd {MONKEY_DIR_PATH} & {CUSTOM_WINDOWS_CMD}" assert pba.command == expected_command @@ -94,8 +84,8 @@ def fake_options_files_only(): @pytest.mark.parametrize("os", [set_os_linux, set_os_windows]) -def test_files_only(fake_options_files_only, os, control_client): - pba = CustomPBA(MagicMock(), control_client) +def test_files_only(fake_options_files_only, os): + pba = CustomPBA(MagicMock(), MagicMock()) pba._set_options(fake_options_files_only) assert pba.command == "" @@ -110,15 +100,15 @@ def fake_options_commands_only(): } -def test_commands_only(fake_options_commands_only, set_os_linux, control_client): - pba = CustomPBA(MagicMock(), control_client) +def test_commands_only(fake_options_commands_only, set_os_linux): + pba = CustomPBA(MagicMock(), MagicMock()) pba._set_options(fake_options_commands_only) assert pba.command == CUSTOM_LINUX_CMD assert pba.filename == "" -def test_commands_only_windows(fake_options_commands_only, set_os_windows, control_client): - pba = CustomPBA(MagicMock(), control_client) +def test_commands_only_windows(fake_options_commands_only, set_os_windows): + pba = CustomPBA(MagicMock(), MagicMock()) pba._set_options(fake_options_commands_only) assert pba.command == CUSTOM_WINDOWS_CMD assert pba.filename == "" diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py b/monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py index ce790f2d9..62e3af713 100644 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py @@ -13,7 +13,6 @@ def spy_send_telemetry(monkeypatch): _spy_send_telemetry.telem_category = None _spy_send_telemetry.data = None - control_client = ControlClient("localhost:5000") - ControlClient.control_client_object = control_client + ControlClient.control_client_object = MagicMock() ControlClient.control_client_object.send_telemetry = MagicMock(side_effect=_spy_send_telemetry) return _spy_send_telemetry