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.
This commit is contained in:
Mike Salvatore 2021-06-28 06:52:33 -04:00
parent fadd978050
commit 85c91f55bb
3 changed files with 25 additions and 17 deletions

View File

@ -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()

View File

@ -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.

View File

@ -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