From bedc8d4f842687d60a9c4f7708c966f1262be53f Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 29 Nov 2021 18:11:06 +0530 Subject: [PATCH 1/6] Agent: Add cleanup logic for ransomware payload --- .../ransomware/ransomware_payload.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 60cdeff84..897d12913 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -1,4 +1,5 @@ import logging +import os from pathlib import Path from typing import Callable, List @@ -26,6 +27,8 @@ class RansomwarePayload: self._leave_readme = leave_readme self._telemetry_messenger = telemetry_messenger + self._readme_incomplete = False + def run_payload(self): if not self._config.target_directory: return @@ -37,7 +40,9 @@ class RansomwarePayload: self._encrypt_files(file_list) if self._config.readme_enabled: + self._readme_incomplete = True self._leave_readme(README_SRC, self._config.target_directory / README_FILE_NAME) + self._readme_incomplete = False def _find_files(self) -> List[Path]: logger.info(f"Collecting files in {self._config.target_directory}") @@ -58,3 +63,18 @@ 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 cleanup(self): + if self._readme_incomplete: + logger.info( + "README.txt file dropping was interrupted. Removing corrupt file and " + "trying again." + ) + try: + os.remove(self._config.target_directory / README_FILE_NAME) + self._leave_readme(README_SRC, self._config.target_directory / README_FILE_NAME) + except Exception as ex: + logger.info( + f"An exception occurred: {str(ex)}. README.txt file dropping was " + "unsuccessful." + ) From f87802678be80833d4e12d2e2a18f422597d8671 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 30 Nov 2021 10:03:02 -0500 Subject: [PATCH 2/6] Tests: Use default parameters in build_ransomware_payload() fixture This allows ransomware payloads with different mocks to be built on a per-test basis with minimal effort and maximal code reuse. --- .../ransomware/test_ransomware_payload.py | 31 ++++++++++--------- 1 file changed, 16 insertions(+), 15 deletions(-) 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..0497e28c6 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 @@ -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() From 14c298e89ce8dbed1ffa1a749cb25e55c0ffb25a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 30 Nov 2021 10:42:12 -0500 Subject: [PATCH 3/6] Agent: Move exception handling from readme_dropper to ransomware_payload --- .../ransomware/ransomware_payload.py | 21 +++++++--- .../ransomware/readme_dropper.py | 10 +---- .../ransomware/test_ransomware_payload.py | 38 ++++++++++++++++++- 3 files changed, 53 insertions(+), 16 deletions(-) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 897d12913..ef5a28289 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -1,5 +1,4 @@ import logging -import os from pathlib import Path from typing import Callable, List @@ -27,6 +26,10 @@ 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): @@ -40,9 +43,7 @@ class RansomwarePayload: self._encrypt_files(file_list) if self._config.readme_enabled: - self._readme_incomplete = True - self._leave_readme(README_SRC, self._config.target_directory / README_FILE_NAME) - self._readme_incomplete = False + self._leave_readme_in_target_directory() def _find_files(self) -> List[Path]: logger.info(f"Collecting files in {self._config.target_directory}") @@ -64,6 +65,14 @@ class RansomwarePayload: 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): if self._readme_incomplete: logger.info( @@ -71,8 +80,8 @@ class RansomwarePayload: "trying again." ) try: - os.remove(self._config.target_directory / README_FILE_NAME) - self._leave_readme(README_SRC, self._config.target_directory / README_FILE_NAME) + self._readme_file_path.unlink() + self._leave_readme_in_target_directory() except Exception as ex: logger.info( f"An exception occurred: {str(ex)}. README.txt file dropping was " 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 0497e28c6..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 @@ -173,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 From 62a6b09e006a3060d4cd93e0286beb31f38d6cfd Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 30 Nov 2021 10:59:18 -0500 Subject: [PATCH 4/6] Agent: Use `self._target_directory` in RansomwarePayload --- monkey/infection_monkey/ransomware/ransomware_payload.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index ef5a28289..afa4ffe25 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -33,7 +33,7 @@ class RansomwarePayload: 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") @@ -46,11 +46,11 @@ class RansomwarePayload: 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: From 789a6691c13425257782d60fac3a4eeb72dceead Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 30 Nov 2021 11:34:45 -0500 Subject: [PATCH 5/6] Agent: Improve log messages in RansomwarePayload.cleanup() --- .../infection_monkey/ransomware/ransomware_payload.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index afa4ffe25..86c7cd9ba 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -76,14 +76,14 @@ class RansomwarePayload: def cleanup(self): if self._readme_incomplete: logger.info( - "README.txt file dropping was interrupted. Removing corrupt file and " - "trying again." + "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.info( - f"An exception occurred: {str(ex)}. README.txt file dropping was " - "unsuccessful." + logger.error( + "An error occurred while trying to remove the corrupt or incomplete README.txt " + f"file: {ex}" ) From a5fc0bc3936bba5f8a3e58b60327f374ea32f0c1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 30 Nov 2021 11:35:04 -0500 Subject: [PATCH 6/6] Agent: Change readme if condition in RansomwarePayload.cleanup() If the _readme_incomplete flag is set but no readme file has been left in the target directory, do not leave a new readme file. This can happen if the thread is forcefully killed between the time when the flag is set and the file is first created. The cleanup function is only concerned with cleaning up incomplete files, not ensuring the existence of the file under all circumstances. --- monkey/infection_monkey/ransomware/ransomware_payload.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 86c7cd9ba..ff2a89d64 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -74,7 +74,10 @@ class RansomwarePayload: logger.warning(f"An error occurred while attempting to leave a README.txt file: {ex}") def cleanup(self): - if self._readme_incomplete: + # 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."