From 85c91f55bb7782253ae203c5cdd4ae96c0c5ec83 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 28 Jun 2021 06:52:33 -0400 Subject: [PATCH] agent: Use BatchingTelemetryMessenger in RansomewarePayload We don't want the ransomware payload to encrypt all files and then send telemetry to the island. This could lead to a long period of time where the user has no insight into what the monkey is doing on a node. We also don't want to flood the island with telemetries. By using the BatchingTelemetryMessenger, ransomware encryption telemetries are batched together and periodically sent to the island. --- monkey/infection_monkey/monkey.py | 10 ++++++- ...ger.py => batching_telemetry_messenger.py} | 2 +- ...y => test_batching_telemetry_messenger.py} | 30 +++++++++---------- 3 files changed, 25 insertions(+), 17 deletions(-) rename monkey/infection_monkey/telemetry/messengers/{batched_telemetry_messenger.py => batching_telemetry_messenger.py} (98%) rename monkey/tests/unit_tests/infection_monkey/telemetry/messengers/{test_batched_telemetry_messenger.py => test_batching_telemetry_messenger.py} (74%) 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