Agent: Accept a "select_files" Callable

monkey/infection_monkey/monkey.py

@@ -6,6 +6,8 @@ import sys
 import time
 from threading import Thread
 
+from InfectionMonkey.ransomware.targeted_file_extensions import TARGETED_FILE_EXTENSIONS
+
 import infection_monkey.tunnel as tunnel
 from common.utils.attack_utils import ScanStatus, UsageEnum
 from common.utils.exceptions import ExploitingVulnerableMachineError, FailedExploitationError
@@ -19,7 +21,8 @@ from infection_monkey.network.HostFinger import HostFinger
 from infection_monkey.network.network_scanner import NetworkScanner
 from infection_monkey.network.tools import get_interface_to_target, is_running_on_island
 from infection_monkey.post_breach.post_breach_handler import PostBreach
-from infection_monkey.ransomware import readme_utils
+from infection_monkey.ransomware import ransomware_payload, readme_utils
+from infection_monkey.ransomware.file_selectors import ProductionSafeTargetFileSelector
 from infection_monkey.ransomware.ransomware_payload import RansomwarePayload
 from infection_monkey.system_info import SystemInfoCollector
 from infection_monkey.system_singleton import SystemSingleton
@@ -476,9 +479,14 @@ class InfectionMonkey(object):
         telemetry_messenger = LegacyTelemetryMessengerAdapter()
         batching_telemetry_messenger = BatchingTelemetryMessenger(telemetry_messenger)
 
+        targeted_file_extensions = TARGETED_FILE_EXTENSIONS.copy()
+        targeted_file_extensions.discard(ransomware_payload.EXTENSION)
+        file_selector = ProductionSafeTargetFileSelector(targeted_file_extensions)
+
         try:
             RansomwarePayload(
                 WormConfiguration.ransomware,
+                file_selector,
                 readme_utils.leave_readme,
                 batching_telemetry_messenger,
             ).run_payload()

monkey/infection_monkey/ransomware/file_selectors.py

@@ -10,12 +10,16 @@ from infection_monkey.utils.dir_utils import (
 )
 
 
-def select_production_safe_target_files(target_dir: Path, extensions: Set) -> List[Path]:
-    file_filters = [
-        file_extension_filter(extensions),
-        is_not_shortcut_filter,
-        is_not_symlink_filter,
-    ]
+class ProductionSafeTargetFileSelector:
+    def __init__(self, targeted_file_extensions: Set[str]):
+        self._targeted_file_extensions = targeted_file_extensions
 
-    all_files = get_all_regular_files_in_directory(target_dir)
-    return filter_files(all_files, file_filters)
+    def __call__(self, target_dir: Path) -> List[Path]:
+        file_filters = [
+            file_extension_filter(self._targeted_file_extensions),
+            is_not_shortcut_filter,
+            is_not_symlink_filter,
+        ]
+
+        all_files = get_all_regular_files_in_directory(target_dir)
+        return filter_files(all_files, file_filters)

monkey/infection_monkey/ransomware/ransomware_payload.py

@@ -5,8 +5,6 @@ from typing import Callable, List, Optional, Tuple
 
 from common.utils.file_utils import InvalidPath, expand_path
 from infection_monkey.ransomware.bitflip_encryptor import BitflipEncryptor
-from infection_monkey.ransomware.file_selectors import select_production_safe_target_files
-from infection_monkey.ransomware.targeted_file_extensions import TARGETED_FILE_EXTENSIONS
 from infection_monkey.telemetry.file_encryption_telem import FileEncryptionTelem
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.utils.environment import is_windows_os
@@ -24,6 +22,7 @@ class RansomwarePayload:
     def __init__(
         self,
         config: dict,
+        select_files: Callable[[Path], List[Path]],
         leave_readme: Callable[[Path, Path], None],
         telemetry_messenger: ITelemetryMessenger,
     ):
@@ -34,10 +33,9 @@ class RansomwarePayload:
 
         self._target_dir = RansomwarePayload.get_target_dir(config)
         self._new_file_extension = EXTENSION
-        self._targeted_file_extensions = TARGETED_FILE_EXTENSIONS.copy()
-        self._targeted_file_extensions.discard(self._new_file_extension)
 
         self._encryptor = BitflipEncryptor(chunk_size=CHUNK_SIZE)
+        self._select_files = select_files
         self._leave_readme = leave_readme
         self._telemetry_messenger = telemetry_messenger
 
@@ -70,9 +68,7 @@ class RansomwarePayload:
 
     def _find_files(self) -> List[Path]:
         LOG.info(f"Collecting files in {self._target_dir}")
-        return sorted(
-            select_production_safe_target_files(self._target_dir, self._targeted_file_extensions)
-        )
+        return sorted(self._select_files(self._target_dir))
 
     def _encrypt_files(self, file_list: List[Path]) -> List[Tuple[Path, Optional[Exception]]]:
         LOG.info(f"Encrypting files in {self._target_dir}")

monkey/tests/unit_tests/infection_monkey/ransomware/conftest.py

@@ -12,8 +12,12 @@ def patched_home_env(monkeypatch, tmp_path):
 
 
 @pytest.fixture
-def ransomware_target(tmp_path, data_for_tests_dir):
-    ransomware_test_data = Path(data_for_tests_dir) / "ransomware_targets"
+def ransomware_test_data(data_for_tests_dir):
+    return Path(data_for_tests_dir) / "ransomware_targets"
+
+
+@pytest.fixture
+def ransomware_target(tmp_path, ransomware_test_data):
     ransomware_target = tmp_path / "ransomware_target"
     shutil.copytree(ransomware_test_data, ransomware_target)
 

monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py

@@ -0,0 +1,56 @@
+import os
+
+import pytest
+from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import (
+    ALL_ZEROS_PDF,
+    HELLO_TXT,
+    SHORTCUT_LNK,
+    SUBDIR,
+    TEST_KEYBOARD_TXT,
+    TEST_LIB_DLL,
+)
+from tests.utils import is_user_admin
+
+from infection_monkey.ransomware.file_selectors import ProductionSafeTargetFileSelector
+
+TARGETED_FILE_EXTENSIONS = [".pdf", ".txt"]
+
+
+@pytest.fixture
+def file_selector():
+    return ProductionSafeTargetFileSelector(TARGETED_FILE_EXTENSIONS)
+
+
+def test_select_targeted_files_only(ransomware_test_data, file_selector):
+    selected_files = file_selector(ransomware_test_data)
+    print(ransomware_test_data)
+
+    assert len(selected_files) == 2
+    assert (ransomware_test_data / ALL_ZEROS_PDF) in selected_files
+    assert (ransomware_test_data / TEST_KEYBOARD_TXT) in selected_files
+
+
+def test_shortcut_not_selected(ransomware_test_data):
+    extensions = TARGETED_FILE_EXTENSIONS + [".lnk"]
+    file_selector = ProductionSafeTargetFileSelector(extensions)
+
+    selected_files = file_selector(ransomware_test_data)
+    assert ransomware_test_data / SHORTCUT_LNK not in selected_files
+
+
+@pytest.mark.skipif(
+    os.name == "nt" and not is_user_admin(), reason="Test requires admin rights on Windows"
+)
+def test_symlink_not_selected(ransomware_target, file_selector):
+    SYMLINK = "symlink.pdf"
+    link_path = ransomware_target / SYMLINK
+    link_path.symlink_to(ransomware_target / TEST_LIB_DLL)
+
+    selected_files = file_selector(ransomware_target)
+    assert link_path not in selected_files
+
+
+def test_directories_not_selected(ransomware_test_data, file_selector):
+    selected_files = file_selector(ransomware_test_data)
+
+    assert (ransomware_test_data / SUBDIR / HELLO_TXT) not in selected_files

monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py

@@ -22,13 +22,14 @@ from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import
 )
 from tests.utils import hash_file, is_user_admin
 
-from infection_monkey.ransomware import ransomware_payload as ransomware_payload_module
+from infection_monkey.ransomware.file_selectors import ProductionSafeTargetFileSelector
 from infection_monkey.ransomware.ransomware_payload import (
     EXTENSION,
     README_DEST,
     README_SRC,
     RansomwarePayload,
 )
+from infection_monkey.ransomware.targeted_file_extensions import TARGETED_FILE_EXTENSIONS
 
 
 def with_extension(filename):
@@ -55,13 +56,20 @@ def ransomware_payload(build_ransomware_payload, ransomware_payload_config):
 
 
 @pytest.fixture
-def build_ransomware_payload(telemetry_messenger_spy, mock_leave_readme):
+def build_ransomware_payload(telemetry_messenger_spy, mock_file_selector, mock_leave_readme):
     def inner(config):
-        return RansomwarePayload(config, mock_leave_readme, telemetry_messenger_spy)
+        return RansomwarePayload(
+            config, mock_file_selector, mock_leave_readme, telemetry_messenger_spy
+        )
 
     return inner
 
 
+@pytest.fixture
+def mock_file_selector():
+    return ProductionSafeTargetFileSelector(TARGETED_FILE_EXTENSIONS)
+
+
 @pytest.fixture
 def mock_leave_readme():
     return MagicMock()
@@ -209,10 +217,12 @@ def test_telemetry_success(ransomware_payload, telemetry_messenger_spy):
     assert telem_2.get_data()["files"][0]["error"] == ""
 
 
-def test_telemetry_failure(monkeypatch, ransomware_payload, telemetry_messenger_spy):
+def test_telemetry_failure(
+    monkeypatch, mock_file_selector, ransomware_payload, telemetry_messenger_spy
+):
     monkeypatch.setattr(
-        ransomware_payload_module,
-        "select_production_safe_target_files",
+        ProductionSafeTargetFileSelector,
+        "__call__",
         lambda a, b: [PurePosixPath("/file/not/exist")],
     ),
 
@@ -251,14 +261,12 @@ def test_no_readme_if_no_directory(
     telemetry_messenger_spy,
     ransomware_target,
 ):
-    monkeypatch.setattr(ransomware_payload_module, "TARGETED_FILE_EXTENSIONS", set()),
-
     ransomware_payload_config["encryption"]["directories"]["linux_target_dir"] = ""
     ransomware_payload_config["encryption"]["directories"]["windows_target_dir"] = ""
     ransomware_payload_config["other_behaviors"]["readme"] = True
 
     RansomwarePayload(
-        ransomware_payload_config, mock_leave_readme, telemetry_messenger_spy
+        ransomware_payload_config, mock_file_selector, mock_leave_readme, telemetry_messenger_spy
     ).run_payload()
 
     mock_leave_readme.assert_not_called()