Agent: Add file extension to ransomware

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_encryptor = _build_file_encryptor(ransomware_options.file_extension)
-    file_selector = _build_file_selector()
+    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)
 

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"

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