agent: Add explicit start to 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. Therefore, a
stop() was added to allow the BatchingTelemetryMessenger to be stopped.
Since it has an explicit stop, it should also have an explicit start,
rather than starting the thread in the constructor.
This commit is contained in:
Mike Salvatore 2021-06-28 12:05:57 -04:00
parent 0a9c98f061
commit 1d066c8e6d
3 changed files with 7 additions and 3 deletions

View File

@ -474,6 +474,7 @@ class InfectionMonkey(object):
def run_ransomware():
telemetry_messenger = LegacyTelemetryMessengerAdapter()
batching_telemetry_messenger = BatchingTelemetryMessenger(telemetry_messenger)
batching_telemetry_messenger.start()
try:
RansomewarePayload(
WormConfiguration.ransomware, batching_telemetry_messenger

View File

@ -28,14 +28,16 @@ class BatchingTelemetryMessenger(ITelemetryMessenger):
self._last_sent_time = time.time()
self._telemetry_batches: Dict[str, IBatchableTelem] = {}
def __del__(self):
self.stop()
def start(self):
self._should_run_batch_thread = True
self._manage_telemetry_batches_thread = threading.Thread(
target=self._manage_telemetry_batches
)
self._manage_telemetry_batches_thread.start()
def __del__(self):
self.stop()
def stop(self):
self._should_run_batch_thread = False
self._manage_telemetry_batches_thread.join()

View File

@ -57,6 +57,7 @@ class BatchableTelemStub(BatchableTelemMixin, BaseTelem, IBatchableTelem):
def batching_telemetry_messenger(monkeypatch, telemetry_messenger_spy):
patch_time(monkeypatch, 0)
btm = BatchingTelemetryMessenger(telemetry_messenger_spy, period=0.001)
btm.start()
yield btm
btm.stop()