diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_builder.py b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py index e52f52cab..c1da93879 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware_builder.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py @@ -1,6 +1,9 @@ import logging from pprint import pformat +from infection_monkey.telemetry.messengers.batching_telemetry_messenger import ( + DEFAULT_PERIOD as DEFAULT_TELEMETRY_BATCH_PERIOD, +) from infection_monkey.telemetry.messengers.batching_telemetry_messenger import ( BatchingTelemetryMessenger, ) @@ -19,7 +22,12 @@ CHUNK_SIZE = 4096 * 24 logger = logging.getLogger(__name__) -def build_ransomware(options: dict, telemetry_messenger: ITelemetryMessenger): +def build_ransomware( + options: dict, + telemetry_messenger: ITelemetryMessenger, + # BatchingTelemetryMessenger will go away soon and so will this parameter + telemetry_batch_period: float = DEFAULT_TELEMETRY_BATCH_PERIOD, +): logger.debug(f"Ransomware configuration:\n{pformat(options)}") ransomware_options = RansomwareOptions(options) @@ -32,7 +40,7 @@ def build_ransomware(options: dict, telemetry_messenger: ITelemetryMessenger): file_encryptor, file_selector, leave_readme, - BatchingTelemetryMessenger(telemetry_messenger), + BatchingTelemetryMessenger(telemetry_messenger, telemetry_batch_period), ) 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 5dcecf840..ced8cf9fd 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 @@ -23,7 +23,9 @@ def test_uses_correct_extension(ransomware_options_dict, tmp_path, ransomware_fi ransomware_directories["linux_target_dir"] = target_dir ransomware_directories["windows_target_dir"] = target_dir telemetry_messenger = MagicMock() - ransomware = ransomware_builder.build_ransomware(ransomware_options_dict, telemetry_messenger) + ransomware = ransomware_builder.build_ransomware( + ransomware_options_dict, telemetry_messenger, 0.01 + ) file = target_dir / "file.txt" file.write_text("Do your worst!")