forked from p15670423/monkey
Agent: Add control_client_object to ControlClient
* Workaround global class attribute
This commit is contained in:
parent
94dbd9a8e2
commit
c467dde145
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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):
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue