p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

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.
This commit is contained in:
Mike Salvatore 2022-09-16 10:34:31 -04:00
parent f6270400b6
commit f677096439
2 changed files with 13 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
monkey/infection_monkey/payload/ransomware/ransomware_builder.py
View File

@ -1,6 +1,9 @@
import logging import logging
from pprint import pformat 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 ( from infection_monkey.telemetry.messengers.batching_telemetry_messenger import (
BatchingTelemetryMessenger, BatchingTelemetryMessenger,
) )
@ -19,7 +22,12 @@ CHUNK_SIZE = 4096 * 24
logger = logging.getLogger(__name__) 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)}") logger.debug(f"Ransomware configuration:\n{pformat(options)}")
ransomware_options = RansomwareOptions(options) ransomware_options = RansomwareOptions(options)
@ -32,7 +40,7 @@ def build_ransomware(options: dict, telemetry_messenger: ITelemetryMessenger):
file_encryptor, file_encryptor,
file_selector, file_selector,
leave_readme, leave_readme,
BatchingTelemetryMessenger(telemetry_messenger), BatchingTelemetryMessenger(telemetry_messenger, telemetry_batch_period),
) )

4
monkey/tests/integration_tests/infection_monkey/payload/ransomware/test_integrated_ransomware.py
View File

@ -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["linux_target_dir"] = target_dir
ransomware_directories["windows_target_dir"] = target_dir ransomware_directories["windows_target_dir"] = target_dir
telemetry_messenger = MagicMock() 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 = target_dir / "file.txt"
file.write_text("Do your worst!") file.write_text("Do your worst!")