diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index dfcf8d327..98b2cab78 100644 --- a/monkey/infection_monkey/monkey.py +++ b/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 diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_builder.py b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py index e2f3b87eb..1a2649060 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware_builder.py +++ b/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) diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_payload.py b/monkey/infection_monkey/payload/ransomware/ransomware_payload.py index d785859a2..79ab5882d 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware_payload.py +++ b/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) diff --git a/monkey/tests/integration_tests/infection_monkey/payload/ransomware/test_integrated_ransomware.py b/monkey/tests/integration_tests/infection_monkey/payload/ransomware/test_integrated_ransomware.py index 973705f23..5dcecf840 100644 --- a/monkey/tests/integration_tests/infection_monkey/payload/ransomware/test_integrated_ransomware.py +++ b/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!")