diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 60cdeff84..ff2a89d64 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -26,8 +26,14 @@ class RansomwarePayload: self._leave_readme = leave_readme self._telemetry_messenger = telemetry_messenger + self._target_directory = self._config.target_directory + self._readme_file_path = ( + self._target_directory / README_FILE_NAME if self._target_directory else None + ) + self._readme_incomplete = False + def run_payload(self): - if not self._config.target_directory: + if not self._target_directory: return logger.info("Running ransomware payload") @@ -37,14 +43,14 @@ class RansomwarePayload: self._encrypt_files(file_list) if self._config.readme_enabled: - self._leave_readme(README_SRC, self._config.target_directory / README_FILE_NAME) + self._leave_readme_in_target_directory() def _find_files(self) -> List[Path]: - logger.info(f"Collecting files in {self._config.target_directory}") - return sorted(self._select_files(self._config.target_directory)) + logger.info(f"Collecting files in {self._target_directory}") + return sorted(self._select_files(self._target_directory)) def _encrypt_files(self, file_list: List[Path]): - logger.info(f"Encrypting files in {self._config.target_directory}") + logger.info(f"Encrypting files in {self._target_directory}") for filepath in file_list: try: @@ -58,3 +64,29 @@ class RansomwarePayload: def _send_telemetry(self, filepath: Path, success: bool, error: str): encryption_attempt = FileEncryptionTelem(str(filepath), success, error) self._telemetry_messenger.send_telemetry(encryption_attempt) + + def _leave_readme_in_target_directory(self): + try: + self._readme_incomplete = True + self._leave_readme(README_SRC, self._readme_file_path) + self._readme_incomplete = False + except Exception as ex: + logger.warning(f"An error occurred while attempting to leave a README.txt file: {ex}") + + def cleanup(self): + # This cleanup function is only concerned with cleaning up and replacing *incomplete* + # README.txt files; its goal is not to ensure the existence of a README file. Therefore, + # only retry if a README.txt file actually exists. + if self._readme_incomplete and self._readme_file_path.exists(): + logger.info( + "The process of leaving a README.txt was interrupted. Removing the corrupt file " + "and trying again." + ) + try: + self._readme_file_path.unlink() + self._leave_readme_in_target_directory() + except Exception as ex: + logger.error( + "An error occurred while trying to remove the corrupt or incomplete README.txt " + f"file: {ex}" + ) diff --git a/monkey/infection_monkey/ransomware/readme_dropper.py b/monkey/infection_monkey/ransomware/readme_dropper.py index 12a171c5b..253c5e574 100644 --- a/monkey/infection_monkey/ransomware/readme_dropper.py +++ b/monkey/infection_monkey/ransomware/readme_dropper.py @@ -10,13 +10,5 @@ def leave_readme(src: Path, dest: Path): logger.warning(f"{dest} already exists, not leaving a new README.txt") return - _copy_readme_file(src, dest) - - -def _copy_readme_file(src: Path, dest: Path): logger.info(f"Leaving a ransomware README file at {dest}") - - try: - shutil.copyfile(src, dest) - except Exception as ex: - logger.warning(f"An error occurred while attempting to leave a README.txt file: {ex}") + shutil.copyfile(src, dest) 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 6c73cfb8d..09a330553 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 @@ -1,4 +1,4 @@ -from pathlib import PurePosixPath +from pathlib import Path, PurePosixPath from unittest.mock import MagicMock import pytest @@ -21,12 +21,17 @@ def ransomware_payload(build_ransomware_payload, ransomware_payload_config): def build_ransomware_payload( mock_file_encryptor, mock_file_selector, mock_leave_readme, telemetry_messenger_spy ): - def inner(config): + def inner( + config, + file_encryptor=mock_file_encryptor, + file_selector=mock_file_selector, + leave_readme=mock_leave_readme, + ): return RansomwarePayload( config, - mock_file_encryptor, - mock_file_selector, - mock_leave_readme, + file_encryptor, + file_selector, + leave_readme, telemetry_messenger_spy, ) @@ -121,19 +126,15 @@ def test_telemetry_success(ransomware_payload, telemetry_messenger_spy): def test_telemetry_failure( - monkeypatch, ransomware_payload_config, mock_leave_readme, telemetry_messenger_spy + build_ransomware_payload, ransomware_payload_config, telemetry_messenger_spy ): file_not_exists = "/file/not/exist" - ransomware_payload = RansomwarePayload( - ransomware_payload_config, - MagicMock( - side_effect=FileNotFoundError( - f"[Errno 2] No such file or directory: '{file_not_exists}'" - ) - ), - MagicMock(return_value=[PurePosixPath(file_not_exists)]), - mock_leave_readme, - telemetry_messenger_spy, + mfe = MagicMock( + side_effect=FileNotFoundError(f"[Errno 2] No such file or directory: '{file_not_exists}'") + ) + mfs = MagicMock(return_value=[PurePosixPath(file_not_exists)]) + ransomware_payload = build_ransomware_payload( + config=ransomware_payload_config, file_encryptor=mfe, file_selector=mfs ) ransomware_payload.run_payload() @@ -172,3 +173,39 @@ def test_no_readme_if_no_directory( ransomware_payload.run_payload() mock_leave_readme.assert_not_called() + + +def test_leave_readme_exceptions_handled(build_ransomware_payload, ransomware_payload_config): + leave_readme = MagicMock(side_effect=Exception("Test exception when leaving README")) + ransomware_payload_config.readme_enabled = True + ransomware_payload = build_ransomware_payload( + config=ransomware_payload_config, leave_readme=leave_readme + ) + + # Test will fail if exception is raised and not handled + ransomware_payload.run_payload() + ransomware_payload.cleanup() + + +def test_cleanup_incomplete_readme(build_ransomware_payload, ransomware_payload_config): + def leave_readme(_: Path, dest: Path): + if leave_readme.i == 0: + dest.touch() + + leave_readme.i += 1 + + raise Exception("Test exception when leaving README") + + leave_readme.i = 0 + + ransomware_payload_config.readme_enabled = True + ransomware_payload = build_ransomware_payload( + config=ransomware_payload_config, leave_readme=leave_readme + ) + + ransomware_payload.run_payload() + assert (ransomware_payload_config.target_directory / README_FILE_NAME).exists() + + ransomware_payload.cleanup() + assert not (ransomware_payload_config.target_directory / README_FILE_NAME).exists() + assert leave_readme.i == 2