forked from p15670423/monkey
agent: Remove start/stop from BatchingTelemetryMessenger
My original plan was to start a thread in __init__() and stop the thread when __del__() was called. Since the running thread (object) contains a reference to the BatchingTelemetryMessenger object that launched it, the destructor will not be called until the thread is stopped. This resulted in adding a stop() method (fadd978
) followed by adding a start() method (1d066c8e
). By using an inner class to run the thread, we enable the class to be used as originally intended, reducing the burden on the user of this class. The thread is now started on construction and stopped on destruction. The user can remain blissfully unaware that anything resembling threading is going in, and can use the BatchingTelemetryMessenger just like any other ITelemetryMessenger.
This commit is contained in:
parent
1d066c8e6d
commit
2f62a14fbf
|
@ -474,12 +474,10 @@ class InfectionMonkey(object):
|
||||||
def run_ransomware():
|
def run_ransomware():
|
||||||
telemetry_messenger = LegacyTelemetryMessengerAdapter()
|
telemetry_messenger = LegacyTelemetryMessengerAdapter()
|
||||||
batching_telemetry_messenger = BatchingTelemetryMessenger(telemetry_messenger)
|
batching_telemetry_messenger = BatchingTelemetryMessenger(telemetry_messenger)
|
||||||
batching_telemetry_messenger.start()
|
|
||||||
try:
|
try:
|
||||||
RansomewarePayload(
|
RansomewarePayload(
|
||||||
WormConfiguration.ransomware, batching_telemetry_messenger
|
WormConfiguration.ransomware, batching_telemetry_messenger
|
||||||
).run_payload()
|
).run_payload()
|
||||||
except Exception as ex:
|
except Exception as ex:
|
||||||
LOG.error(f"An unexpected error occurred while running the ransomware payload: {ex}")
|
LOG.error(f"An unexpected error occurred while running the ransomware payload: {ex}")
|
||||||
finally:
|
|
||||||
batching_telemetry_messenger.stop()
|
|
||||||
|
|
|
@ -18,19 +18,31 @@ class BatchingTelemetryMessenger(ITelemetryMessenger):
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, telemetry_messenger: ITelemetryMessenger, period=DEFAULT_PERIOD):
|
def __init__(self, telemetry_messenger: ITelemetryMessenger, period=DEFAULT_PERIOD):
|
||||||
|
self._queue: queue.Queue[ITelem] = queue.Queue()
|
||||||
|
self._thread = BatchingTelemetryMessenger._BatchingTelemetryMessengerThread(
|
||||||
|
self._queue, telemetry_messenger, period
|
||||||
|
)
|
||||||
|
|
||||||
|
self._thread.start()
|
||||||
|
|
||||||
|
def __del__(self):
|
||||||
|
self._thread.stop()
|
||||||
|
|
||||||
|
def send_telemetry(self, telemetry: ITelem):
|
||||||
|
self._queue.put(telemetry)
|
||||||
|
|
||||||
|
class _BatchingTelemetryMessengerThread:
|
||||||
|
def __init__(self, queue: queue.Queue, telemetry_messenger: ITelemetryMessenger, period):
|
||||||
|
self._queue: queue.Queue[ITelem] = queue
|
||||||
self._telemetry_messenger = telemetry_messenger
|
self._telemetry_messenger = telemetry_messenger
|
||||||
self._period = period
|
self._period = period
|
||||||
|
|
||||||
self._should_run_batch_thread = True
|
self._should_run_batch_thread = True
|
||||||
self._queue: queue.Queue[ITelem] = queue.Queue()
|
|
||||||
# TODO: Create a "timer" or "countdown" class and inject an object instead of
|
# TODO: Create a "timer" or "countdown" class and inject an object instead of
|
||||||
# using time.time()
|
# using time.time()
|
||||||
self._last_sent_time = time.time()
|
self._last_sent_time = time.time()
|
||||||
self._telemetry_batches: Dict[str, IBatchableTelem] = {}
|
self._telemetry_batches: Dict[str, IBatchableTelem] = {}
|
||||||
|
|
||||||
def __del__(self):
|
|
||||||
self.stop()
|
|
||||||
|
|
||||||
def start(self):
|
def start(self):
|
||||||
self._should_run_batch_thread = True
|
self._should_run_batch_thread = True
|
||||||
self._manage_telemetry_batches_thread = threading.Thread(
|
self._manage_telemetry_batches_thread = threading.Thread(
|
||||||
|
@ -42,9 +54,6 @@ class BatchingTelemetryMessenger(ITelemetryMessenger):
|
||||||
self._should_run_batch_thread = False
|
self._should_run_batch_thread = False
|
||||||
self._manage_telemetry_batches_thread.join()
|
self._manage_telemetry_batches_thread.join()
|
||||||
|
|
||||||
def send_telemetry(self, telemetry: ITelem):
|
|
||||||
self._queue.put(telemetry)
|
|
||||||
|
|
||||||
def _manage_telemetry_batches(self):
|
def _manage_telemetry_batches(self):
|
||||||
self._reset()
|
self._reset()
|
||||||
|
|
||||||
|
|
|
@ -56,11 +56,7 @@ class BatchableTelemStub(BatchableTelemMixin, BaseTelem, IBatchableTelem):
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def batching_telemetry_messenger(monkeypatch, telemetry_messenger_spy):
|
def batching_telemetry_messenger(monkeypatch, telemetry_messenger_spy):
|
||||||
patch_time(monkeypatch, 0)
|
patch_time(monkeypatch, 0)
|
||||||
btm = BatchingTelemetryMessenger(telemetry_messenger_spy, period=0.001)
|
return BatchingTelemetryMessenger(telemetry_messenger_spy, period=0.001)
|
||||||
btm.start()
|
|
||||||
yield btm
|
|
||||||
|
|
||||||
btm.stop()
|
|
||||||
|
|
||||||
|
|
||||||
def test_send_immediately(batching_telemetry_messenger, telemetry_messenger_spy):
|
def test_send_immediately(batching_telemetry_messenger, telemetry_messenger_spy):
|
||||||
|
|
Loading…
Reference in New Issue