Agent: Pass ITelemetryMessenger to Ransomware

monkey/infection_monkey/monkey.py

@@ -66,6 +66,9 @@ from infection_monkey.puppet.puppet import Puppet
 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.messengers.batching_telemetry_messenger import (
+    BatchingTelemetryMessenger,
+)
 from infection_monkey.telemetry.messengers.exploit_intercepting_telemetry_messenger import (
     ExploitInterceptingTelemetryMessenger,
 )
@@ -354,7 +357,11 @@ class InfectionMonkey:
             PluginType.POST_BREACH_ACTION,
         )
 
-        puppet.load_plugin("ransomware", RansomwarePayload(), PluginType.PAYLOAD)
+        puppet.load_plugin(
+            "ransomware",
+            RansomwarePayload(BatchingTelemetryMessenger(self._telemetry_messenger)),
+            PluginType.PAYLOAD,
+        )
 
         return puppet
 

monkey/infection_monkey/payload/ransomware/ransomware_builder.py

@@ -1,12 +1,7 @@
 import logging
 from pprint import pformat
 
-from infection_monkey.telemetry.messengers.batching_telemetry_messenger import (
-    BatchingTelemetryMessenger,
-)
-from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter import (
-    LegacyTelemetryMessengerAdapter,
-)
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.utils.bit_manipulators import flip_bits
 
 from . import readme_dropper
@@ -21,14 +16,13 @@ CHUNK_SIZE = 4096 * 24
 logger = logging.getLogger(__name__)
 
 
-def build_ransomware(options: dict):
+def build_ransomware(options: dict, telemetry_messenger: ITelemetryMessenger):
     logger.debug(f"Ransomware configuration:\n{pformat(options)}")
     ransomware_options = RansomwareOptions(options)
 
     file_encryptor = _build_file_encryptor(ransomware_options.file_extension)
     file_selector = _build_file_selector(ransomware_options.file_extension)
     leave_readme = _build_leave_readme()
-    telemetry_messenger = _build_telemetry_messenger()
 
     return Ransomware(
         ransomware_options,
@@ -55,9 +49,3 @@ def _build_file_selector(file_extension: str):
 
 def _build_leave_readme():
     return readme_dropper.leave_readme
-
-
-def _build_telemetry_messenger():
-    telemetry_messenger = LegacyTelemetryMessengerAdapter()
-
-    return BatchingTelemetryMessenger(telemetry_messenger)

monkey/infection_monkey/payload/ransomware/ransomware_payload.py

@@ -2,11 +2,15 @@ import threading
 from typing import Dict
 
 from infection_monkey.payload.i_payload import IPayload
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
 from . import ransomware_builder
 
 
 class RansomwarePayload(IPayload):
+    def __init__(self, telemetry_messenger: ITelemetryMessenger):
+        self._telemetry_messenger = telemetry_messenger
+
     def run(self, options: Dict, interrupt: threading.Event):
-        ransomware = ransomware_builder.build_ransomware(options)
+        ransomware = ransomware_builder.build_ransomware(options, self._telemetry_messenger)
         ransomware.run(interrupt)

monkey/tests/integration_tests/infection_monkey/payload/ransomware/test_integrated_ransomware.py

@@ -1,4 +1,5 @@
 import threading
+from unittest.mock import MagicMock
 
 import pytest
 
@@ -21,7 +22,8 @@ def test_uses_correct_extension(ransomware_options_dict, tmp_path, ransomware_fi
     ransomware_directories = ransomware_options_dict["encryption"]["directories"]
     ransomware_directories["linux_target_dir"] = target_dir
     ransomware_directories["windows_target_dir"] = target_dir
-    ransomware = ransomware_builder.build_ransomware(ransomware_options_dict)
+    telemetry_messenger = MagicMock()
+    ransomware = ransomware_builder.build_ransomware(ransomware_options_dict, telemetry_messenger)
 
     file = target_dir / "file.txt"
     file.write_text("Do your worst!")