From 81eba6e883e63968af66f2b568e024b0d8a1343f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 13 Jul 2021 19:22:42 -0400 Subject: [PATCH] Agent: Accept a "select_files" Callable --- monkey/infection_monkey/monkey.py | 10 +++- .../ransomware/file_selectors.py | 20 ++++--- .../ransomware/ransomware_payload.py | 10 +--- .../infection_monkey/ransomware/conftest.py | 8 ++- .../ransomware/test_file_selectors.py | 56 +++++++++++++++++++ .../ransomware/test_ransomware_payload.py | 26 ++++++--- 6 files changed, 103 insertions(+), 27 deletions(-) create mode 100644 monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index ca6976f91..0dcbbcd17 100644 --- a/monkey/infection_monkey/monkey.py +++ b/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() diff --git a/monkey/infection_monkey/ransomware/file_selectors.py b/monkey/infection_monkey/ransomware/file_selectors.py index f34bc9ca4..167c547e8 100644 --- a/monkey/infection_monkey/ransomware/file_selectors.py +++ b/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) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index b8f0d8805..7b3a9a42a 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/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}") diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/conftest.py b/monkey/tests/unit_tests/infection_monkey/ransomware/conftest.py index a23751633..1e357c798 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/conftest.py +++ b/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) diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py new file mode 100644 index 000000000..56421be3e --- /dev/null +++ b/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 diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py index ce6b7b08e..b3d9269af 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py +++ b/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()