diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 845f754f7..506bc5db2 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -25,6 +25,9 @@ from infection_monkey.system_singleton import SystemSingleton from infection_monkey.telemetry.attack.t1106_telem import T1106Telem from infection_monkey.telemetry.attack.t1107_telem import T1107Telem from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem +from infection_monkey.telemetry.messengers.batching_telemetry_messenger import ( + BatchingTelemetryMessenger, +) from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter import ( LegacyTelemetryMessengerAdapter, ) @@ -471,6 +474,11 @@ class InfectionMonkey(object): def run_ransomware(): try: telemetry_messenger = LegacyTelemetryMessengerAdapter() - RansomewarePayload(WormConfiguration.ransomware, telemetry_messenger).run_payload() + batching_telemetry_messenger = BatchingTelemetryMessenger(telemetry_messenger) + RansomewarePayload( + WormConfiguration.ransomware, batching_telemetry_messenger + ).run_payload() except Exception as ex: LOG.error(f"An unexpected error occurred while running the ransomware payload: {ex}") + finally: + batching_telemetry_messenger.stop() diff --git a/monkey/infection_monkey/telemetry/messengers/batched_telemetry_messenger.py b/monkey/infection_monkey/telemetry/messengers/batching_telemetry_messenger.py similarity index 98% rename from monkey/infection_monkey/telemetry/messengers/batched_telemetry_messenger.py rename to monkey/infection_monkey/telemetry/messengers/batching_telemetry_messenger.py index 16551ef11..43fbf1306 100644 --- a/monkey/infection_monkey/telemetry/messengers/batched_telemetry_messenger.py +++ b/monkey/infection_monkey/telemetry/messengers/batching_telemetry_messenger.py @@ -10,7 +10,7 @@ from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemet WAKES_PER_PERIOD = 4 -class BatchedTelemetryMessenger(ITelemetryMessenger): +class BatchingTelemetryMessenger(ITelemetryMessenger): """ An ITelemetryMessenger decorator that aggregates IBatchableTelem telemetries and periodically sends them to Monkey Island. diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_batched_telemetry_messenger.py b/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_batching_telemetry_messenger.py similarity index 74% rename from monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_batched_telemetry_messenger.py rename to monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_batching_telemetry_messenger.py index 84377ba7c..58ee96d89 100644 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_batched_telemetry_messenger.py +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_batching_telemetry_messenger.py @@ -5,8 +5,8 @@ import pytest from infection_monkey.telemetry.base_telem import BaseTelem from infection_monkey.telemetry.batchable_telem_mixin import BatchableTelemMixin from infection_monkey.telemetry.i_batchable_telem import IBatchableTelem -from infection_monkey.telemetry.messengers.batched_telemetry_messenger import ( - BatchedTelemetryMessenger, +from infection_monkey.telemetry.messengers.batching_telemetry_messenger import ( + BatchingTelemetryMessenger, ) PERIOD = 0.001 @@ -54,31 +54,31 @@ class BatchableTelemStub(BatchableTelemMixin, BaseTelem, IBatchableTelem): @pytest.fixture -def batched_telemetry_messenger(monkeypatch, telemetry_messenger_spy): +def batching_telemetry_messenger(monkeypatch, telemetry_messenger_spy): patch_time(monkeypatch, 0) - btm = BatchedTelemetryMessenger(telemetry_messenger_spy, period=0.001) + btm = BatchingTelemetryMessenger(telemetry_messenger_spy, period=0.001) yield btm btm.stop() -def test_send_immediately(batched_telemetry_messenger, telemetry_messenger_spy): +def test_send_immediately(batching_telemetry_messenger, telemetry_messenger_spy): telem = NonBatchableTelemStub() - batched_telemetry_messenger.send_telemetry(telem) + batching_telemetry_messenger.send_telemetry(telem) release_GIL() assert len(telemetry_messenger_spy.telemetries) == 1 assert telemetry_messenger_spy.telemetries[0] == telem -def test_send_telem_batch(monkeypatch, batched_telemetry_messenger, telemetry_messenger_spy): +def test_send_telem_batch(monkeypatch, batching_telemetry_messenger, telemetry_messenger_spy): expected_data = {"entries": [1, 2]} telem1 = BatchableTelemStub(1) telem2 = BatchableTelemStub(2) - batched_telemetry_messenger.send_telemetry(telem1) - batched_telemetry_messenger.send_telemetry(telem2) + batching_telemetry_messenger.send_telemetry(telem1) + batching_telemetry_messenger.send_telemetry(telem2) release_GIL() assert len(telemetry_messenger_spy.telemetries) == 0 @@ -90,13 +90,13 @@ def test_send_telem_batch(monkeypatch, batched_telemetry_messenger, telemetry_me def test_send_different_telem_types( - monkeypatch, batched_telemetry_messenger, telemetry_messenger_spy + monkeypatch, batching_telemetry_messenger, telemetry_messenger_spy ): telem1 = BatchableTelemStub(1, "cat1") telem2 = BatchableTelemStub(2, "cat2") - batched_telemetry_messenger.send_telemetry(telem1) - batched_telemetry_messenger.send_telemetry(telem2) + batching_telemetry_messenger.send_telemetry(telem1) + batching_telemetry_messenger.send_telemetry(telem2) release_GIL() assert len(telemetry_messenger_spy.telemetries) == 0 @@ -108,15 +108,15 @@ def test_send_different_telem_types( assert telemetry_messenger_spy.telemetries[1] == telem2 -def test_send_two_batches(monkeypatch, batched_telemetry_messenger, telemetry_messenger_spy): +def test_send_two_batches(monkeypatch, batching_telemetry_messenger, telemetry_messenger_spy): telem1 = BatchableTelemStub(1, "cat1") telem2 = BatchableTelemStub(2, "cat1") - batched_telemetry_messenger.send_telemetry(telem1) + batching_telemetry_messenger.send_telemetry(telem1) advance_clock_to_next_period(monkeypatch) release_GIL() - batched_telemetry_messenger.send_telemetry(telem2) + batching_telemetry_messenger.send_telemetry(telem2) release_GIL() assert len(telemetry_messenger_spy.telemetries) == 1