From f6770964397c6541e2e48e44d6bb63ff24e0b91d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 16 Sep 2022 10:34:31 -0400 Subject: [PATCH] Agent: Make BatchingTelemetryMessenger sleep period configurable The ransomware integration tests were sometimes painfully slow because they were waiting for the BatchingTelemetryMessenger to stop its internal thread. This commit surfaces the sleep period parameter to the `build_ransomware()` function so that the unit tests can set it to a very small value. This is somewhat of a hack, but BatchingTelemetryMessenger is going away very soon, as Telemetry is being replaced by Events. --- .../payload/ransomware/ransomware_builder.py | 12 ++++++++++-- .../payload/ransomware/test_integrated_ransomware.py | 4 +++- 2 files changed, 13 insertions(+), 3 deletions(-) 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!")