p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
monkey/infection_monkey/control.py
View File

@ -25,6 +25,11 @@ PBA_FILE_DOWNLOAD = "https://%s/api/pba/download/%s"
class ControlClient: 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): def __init__(self, server_address: str, proxies: Optional[Mapping[str, str]] = None):
self.proxies = {} if not proxies else proxies self.proxies = {} if not proxies else proxies
self.server_address = server_address self.server_address = server_address

2
monkey/infection_monkey/monkey.py
View File

@ -90,6 +90,8 @@ class InfectionMonkey:
self._opts = self._get_arguments(args) self._opts = self._get_arguments(args)
self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.server) self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.server)
self.cc_client = ControlClient(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._monkey_inbound_tunnel = None
self._telemetry_messenger = LegacyTelemetryMessengerAdapter() self._telemetry_messenger = LegacyTelemetryMessengerAdapter()
self._current_depth = self._opts.depth self._current_depth = self._opts.depth

2
monkey/infection_monkey/telemetry/base_telem.py
View File

@ -35,7 +35,7 @@ class BaseTelem(ITelem, metaclass=abc.ABCMeta):
data = self.get_data() data = self.get_data()
serialized_data = json.dumps(data, cls=self.json_encoder) serialized_data = json.dumps(data, cls=self.json_encoder)
self._log_telem_sending(serialized_data, log_data) 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 @property
def json_encoder(self): def json_encoder(self):

4
monkey/tests/unit_tests/infection_monkey/telemetry/conftest.py
View File

@ -11,5 +11,7 @@ def spy_send_telemetry(monkeypatch):
_spy_send_telemetry.telem_category = None _spy_send_telemetry.telem_category = None
_spy_send_telemetry.data = 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 return _spy_send_telemetry