diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 32530b287..27cd6dca1 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -1,12 +1,10 @@ import logging from pathlib import Path -from pprint import pformat from typing import Callable, List -from common.utils.file_utils import InvalidPath, expand_path +from infection_monkey.ransomware.ransomware_config import RansomwareConfig 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 LOG = logging.getLogger(__name__) @@ -17,57 +15,38 @@ README_DEST = "README.txt" class RansomwarePayload: def __init__( self, - config: dict, + config: RansomwareConfig, encrypt_file: Callable[[Path], None], select_files: Callable[[Path], List[Path]], leave_readme: Callable[[Path, Path], None], telemetry_messenger: ITelemetryMessenger, ): - LOG.debug(f"Ransomware payload configuration:\n{pformat(config)}") - - self._encryption_enabled = config["encryption"]["enabled"] - self._readme_enabled = config["other_behaviors"]["readme"] - - self._target_dir = RansomwarePayload.get_target_dir(config) + self._config = config self._encrypt_file = encrypt_file self._select_files = select_files self._leave_readme = leave_readme self._telemetry_messenger = telemetry_messenger - @staticmethod - def get_target_dir(config: dict): - target_directories = config["encryption"]["directories"] - if is_windows_os(): - target_dir_field = target_directories["windows_target_dir"] - else: - target_dir_field = target_directories["linux_target_dir"] - - try: - return expand_path(target_dir_field) - except InvalidPath as e: - LOG.debug(f"Target ransomware dir set to None: {e}") - return None - def run_payload(self): - if not self._target_dir: + if not self._config.target_directory: return LOG.info("Running ransomware payload") - if self._encryption_enabled: + if self._config.encryption_enabled: file_list = self._find_files() self._encrypt_files(file_list) - if self._readme_enabled: - self._leave_readme(README_SRC, self._target_dir / README_DEST) + if self._config.readme_enabled: + self._leave_readme(README_SRC, self._config.target_directory / README_DEST) def _find_files(self) -> List[Path]: - LOG.info(f"Collecting files in {self._target_dir}") - return sorted(self._select_files(self._target_dir)) + LOG.info(f"Collecting files in {self._config.target_directory}") + return sorted(self._select_files(self._config.target_directory)) def _encrypt_files(self, file_list: List[Path]): - LOG.info(f"Encrypting files in {self._target_dir}") + LOG.info(f"Encrypting files in {self._config.target_directory}") for filepath in file_list: try: diff --git a/monkey/infection_monkey/ransomware/ransomware_payload_builder.py b/monkey/infection_monkey/ransomware/ransomware_payload_builder.py index 28770668d..4b13b4082 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload_builder.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload_builder.py @@ -1,6 +1,10 @@ -from infection_monkey.ransomware import readme_utils +import logging +from pprint import pformat + +from infection_monkey.ransomware import ransomware_payload, readme_utils from infection_monkey.ransomware.file_selectors import ProductionSafeTargetFileSelector from infection_monkey.ransomware.in_place_file_encryptor import InPlaceFileEncryptor +from infection_monkey.ransomware.ransomware_config import RansomwareConfig from infection_monkey.ransomware.ransomware_payload import RansomwarePayload from infection_monkey.ransomware.targeted_file_extensions import TARGETED_FILE_EXTENSIONS from infection_monkey.telemetry.messengers.batching_telemetry_messenger import ( @@ -14,14 +18,23 @@ from infection_monkey.utils.bit_manipulators import flip_bits EXTENSION = ".m0nk3y" CHUNK_SIZE = 4096 * 24 +LOG = logging.getLogger(__name__) + def build_ransomware_payload(config: dict): + LOG.debug(f"Ransomware payload configuration:\n{pformat(config)}") + ransomware_config = RansomwareConfig(config) + file_encryptor = _build_file_encryptor() file_selector = _build_file_selector() telemetry_messenger = _build_telemetry_messenger() return RansomwarePayload( - config, file_encryptor, file_selector, readme_utils.leave_readme, telemetry_messenger + ransomware_config, + file_encryptor, + file_selector, + readme_utils.leave_readme, + telemetry_messenger, ) @@ -33,7 +46,7 @@ def _build_file_encryptor(): def _build_file_selector(): targeted_file_extensions = TARGETED_FILE_EXTENSIONS.copy() - targeted_file_extensions.discard(EXTENSION) + targeted_file_extensions.discard(ransomware_payload.EXTENSION) return ProductionSafeTargetFileSelector(targeted_file_extensions) 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 ca52bef5c..118ba8a27 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 @@ -7,6 +7,7 @@ from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import TEST_KEYBOARD_TXT, ) +from infection_monkey.ransomware.ransomware_config import RansomwareConfig from infection_monkey.ransomware.ransomware_payload import ( README_DEST, README_SRC, @@ -14,20 +15,6 @@ from infection_monkey.ransomware.ransomware_payload import ( ) -@pytest.fixture -def ransomware_payload_config(ransomware_target): - return { - "encryption": { - "enabled": True, - "directories": { - "linux_target_dir": str(ransomware_target), - "windows_target_dir": str(ransomware_target), - }, - }, - "other_behaviors": {"readme": False}, - } - - @pytest.fixture def ransomware_payload(build_ransomware_payload, ransomware_payload_config): return build_ransomware_payload(ransomware_payload_config) @@ -50,15 +37,26 @@ def build_ransomware_payload( @pytest.fixture -def mock_file_encryptor(ransomware_target): +def ransomware_payload_config(ransomware_test_data): + class RansomwareConfigStub(RansomwareConfig): + def __init__(self, encryption_enabled, readme_enabled, target_directory): + self.encryption_enabled = encryption_enabled + self.readme_enabled = readme_enabled + self.target_directory = target_directory + + return RansomwareConfigStub(True, False, ransomware_test_data) + + +@pytest.fixture +def mock_file_encryptor(): return MagicMock() @pytest.fixture -def mock_file_selector(ransomware_target): +def mock_file_selector(ransomware_test_data): selected_files = [ - ransomware_target / ALL_ZEROS_PDF, - ransomware_target / TEST_KEYBOARD_TXT, + ransomware_test_data / ALL_ZEROS_PDF, + ransomware_test_data / TEST_KEYBOARD_TXT, ] return MagicMock(return_value=selected_files) @@ -68,37 +66,29 @@ def mock_leave_readme(): return MagicMock() -def test_env_variables_in_target_dir_resolved_linux( +def test_files_selected_from_target_dir( + ransomware_payload, ransomware_payload_config, - build_ransomware_payload, - ransomware_target, - patched_home_env, mock_file_selector, ): - path_with_env_variable = "$HOME/ransomware_target" - - ransomware_payload_config["encryption"]["directories"][ - "linux_target_dir" - ] = ransomware_payload_config["encryption"]["directories"][ - "windows_target_dir" - ] = path_with_env_variable - build_ransomware_payload(ransomware_payload_config).run_payload() - - mock_file_selector.assert_called_with(ransomware_target) + ransomware_payload.run_payload() + mock_file_selector.assert_called_with(ransomware_payload_config.target_directory) -def test_all_selected_files_encrypted(ransomware_target, ransomware_payload, mock_file_encryptor): +def test_all_selected_files_encrypted( + ransomware_test_data, ransomware_payload, mock_file_encryptor +): ransomware_payload.run_payload() assert mock_file_encryptor.call_count == 2 - mock_file_encryptor.assert_any_call(ransomware_target / ALL_ZEROS_PDF) - mock_file_encryptor.assert_any_call(ransomware_target / TEST_KEYBOARD_TXT) + mock_file_encryptor.assert_any_call(ransomware_test_data / ALL_ZEROS_PDF) + mock_file_encryptor.assert_any_call(ransomware_test_data / TEST_KEYBOARD_TXT) def test_encryption_skipped_if_configured_false( - build_ransomware_payload, ransomware_payload_config, ransomware_target, mock_file_encryptor + build_ransomware_payload, ransomware_payload_config, mock_file_encryptor ): - ransomware_payload_config["encryption"]["enabled"] = False + ransomware_payload_config.encryption_enabled = False ransomware_payload = build_ransomware_payload(ransomware_payload_config) ransomware_payload.run_payload() @@ -109,9 +99,8 @@ def test_encryption_skipped_if_configured_false( def test_encryption_skipped_if_no_directory( build_ransomware_payload, ransomware_payload_config, mock_file_encryptor ): - ransomware_payload_config["encryption"]["enabled"] = True - ransomware_payload_config["encryption"]["directories"]["linux_target_dir"] = "" - ransomware_payload_config["encryption"]["directories"]["windows_target_dir"] = "" + ransomware_payload_config.encryption_enabled = True + ransomware_payload_config.target_directory = None ransomware_payload = build_ransomware_payload(ransomware_payload_config) ransomware_payload.run_payload() @@ -158,10 +147,8 @@ def test_telemetry_failure( assert "No such file or directory" in telem.get_data()["files"][0]["error"] -def test_readme_false( - build_ransomware_payload, ransomware_payload_config, mock_leave_readme, ransomware_target -): - ransomware_payload_config["other_behaviors"]["readme"] = False +def test_readme_false(build_ransomware_payload, ransomware_payload_config, mock_leave_readme): + ransomware_payload_config.readme_enabled = False ransomware_payload = build_ransomware_payload(ransomware_payload_config) ransomware_payload.run_payload() @@ -169,21 +156,20 @@ def test_readme_false( def test_readme_true( - build_ransomware_payload, ransomware_payload_config, mock_leave_readme, ransomware_target + build_ransomware_payload, ransomware_payload_config, mock_leave_readme, ransomware_test_data ): - ransomware_payload_config["other_behaviors"]["readme"] = True + ransomware_payload_config.readme_enabled = True ransomware_payload = build_ransomware_payload(ransomware_payload_config) ransomware_payload.run_payload() - mock_leave_readme.assert_called_with(README_SRC, ransomware_target / README_DEST) + mock_leave_readme.assert_called_with(README_SRC, ransomware_test_data / README_DEST) def test_no_readme_if_no_directory( build_ransomware_payload, ransomware_payload_config, mock_leave_readme ): - ransomware_payload_config["encryption"]["directories"]["linux_target_dir"] = "" - ransomware_payload_config["encryption"]["directories"]["windows_target_dir"] = "" - ransomware_payload_config["other_behaviors"]["readme"] = True + ransomware_payload_config.target_directory = None + ransomware_payload_config.readme_enabled = True ransomware_payload = build_ransomware_payload(ransomware_payload_config)