Agent: Add control_client_object to ControlClient

* Workaround global class attribute
This commit is contained in:
Ilija Lazoroski 2022-06-13 13:58:12 +02:00 committed by vakarisz
parent 94dbd9a8e2
commit c467dde145
4 changed files with 11 additions and 2 deletions

View File

@ -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

View File

@ -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

View File

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

View File

@ -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