From c467dde145faae0da824a93ecfc05908b93c5f5b Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 13 Jun 2022 13:58:12 +0200 Subject: [PATCH] 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