Agent: Add control_client_object to ControlClient

* Workaround global class attribute

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

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

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):

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