diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_builder.py b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py index 4b8bbc8bb..9da6f2b33 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware_builder.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py @@ -26,8 +26,8 @@ def build_ransomware(options: dict): logger.debug(f"Ransomware configuration:\n{pformat(options)}") ransomware_options = RansomwareOptions(options) - file_encryptor = _build_file_encryptor() - file_selector = _build_file_selector() + 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() @@ -40,15 +40,16 @@ def build_ransomware(options: dict): ) -def _build_file_encryptor(): +def _build_file_encryptor(file_extension: str): return InPlaceFileEncryptor( - encrypt_bytes=flip_bits, new_file_extension=EXTENSION, chunk_size=CHUNK_SIZE + encrypt_bytes=flip_bits, new_file_extension=file_extension, chunk_size=CHUNK_SIZE ) -def _build_file_selector(): +def _build_file_selector(file_extension: str): targeted_file_extensions = TARGETED_FILE_EXTENSIONS.copy() - targeted_file_extensions.discard(EXTENSION) + if file_extension: + targeted_file_extensions.discard(EXTENSION) return ProductionSafeTargetFileSelector(targeted_file_extensions) diff --git a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/conftest.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/conftest.py index 1e357c798..88f6f60aa 100644 --- a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/conftest.py +++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/conftest.py @@ -11,6 +11,11 @@ def patched_home_env(monkeypatch, tmp_path): return tmp_path +@pytest.fixture(params=[".m0nk3y", ".test", ""], ids=["monkeyext", "testext", "noext"]) +def ransomware_file_extension(request): + return request.param + + @pytest.fixture def ransomware_test_data(data_for_tests_dir): return Path(data_for_tests_dir) / "ransomware_targets" diff --git a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py index 88f37037c..b651dd012 100644 --- a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py +++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py @@ -41,14 +41,15 @@ def build_ransomware( @pytest.fixture -def ransomware_options(ransomware_test_data): +def ransomware_options(ransomware_file_extension, ransomware_test_data): class RansomwareOptionsStub(RansomwareOptions): - def __init__(self, encryption_enabled, readme_enabled, target_directory): + def __init__(self, encryption_enabled, readme_enabled, file_extension, target_directory): self.encryption_enabled = encryption_enabled self.readme_enabled = readme_enabled + self.file_extension = file_extension self.target_directory = target_directory - return RansomwareOptionsStub(True, False, ransomware_test_data) + return RansomwareOptionsStub(True, False, ransomware_file_extension, ransomware_test_data) @pytest.fixture