From f0e9109f6417a5c528d7ba9a5901b819b2ff0fc2 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 8 Jul 2021 11:23:15 -0400 Subject: [PATCH 1/4] Agent: Inject copy_file callable into RansomwarePayload In order to test certain conditions, our options are to either monkeypatch shutil.copyfile(), or inject a callable into the RansomwarePayload. Monkeypatching shutil.copyfile() could lead to issues down the road. For example, if the implementation of `_leave_readme()` is changed to no longer use copyfile(), a test that asserts that copyfile() has not been called will pass, even though a file may have been copied. --- monkey/infection_monkey/monkey.py | 3 +- .../ransomware/ransomware_payload.py | 13 +++++--- .../ransomware/test_ransomware_payload.py | 33 ++++++++++++------- 3 files changed, 32 insertions(+), 17 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index a75304660..e89b9ab2c 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -1,6 +1,7 @@ import argparse import logging import os +import shutil import subprocess import sys import time @@ -477,7 +478,7 @@ class InfectionMonkey(object): try: RansomwarePayload( - WormConfiguration.ransomware, batching_telemetry_messenger + WormConfiguration.ransomware, batching_telemetry_messenger, shutil.copyfile ).run_payload() except Exception as ex: LOG.error(f"An unexpected error occurred while running the ransomware payload: {ex}") diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 29d1fcd65..05c2159d2 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -1,8 +1,7 @@ import logging -import shutil from pathlib import Path from pprint import pformat -from typing import List, Optional, Tuple +from typing import Callable, List, Optional, Tuple from common.utils.file_utils import InvalidPath, expand_path from infection_monkey.ransomware.bitflip_encryptor import BitflipEncryptor @@ -22,7 +21,12 @@ README_DEST = "README.txt" class RansomwarePayload: - def __init__(self, config: dict, telemetry_messenger: ITelemetryMessenger): + def __init__( + self, + config: dict, + telemetry_messenger: ITelemetryMessenger, + copy_file: Callable[[str, str], None], + ): LOG.debug(f"Ransomware payload configuration:\n{pformat(config)}") self._encryption_enabled = config["encryption"]["enabled"] @@ -34,6 +38,7 @@ class RansomwarePayload: self._valid_file_extensions_for_encryption.discard(self._new_file_extension) self._encryptor = BitflipEncryptor(chunk_size=CHUNK_SIZE) + self._copy_file = copy_file self._telemetry_messenger = telemetry_messenger @staticmethod @@ -97,6 +102,6 @@ class RansomwarePayload: LOG.info(f"Leaving a ransomware README file at {readme_dest_path}") try: - shutil.copyfile(README_SRC, readme_dest_path) + self._copy_file(README_SRC, readme_dest_path) except Exception as ex: LOG.warning(f"An error occurred while attempting to leave a README.txt file: {ex}") 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 6a37ee6e5..89c9f5ff5 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,5 @@ import os +import shutil from pathlib import Path, PurePosixPath import pytest @@ -44,12 +45,20 @@ def ransomware_payload_config(ransomware_target): @pytest.fixture -def ransomware_payload(ransomware_payload_config, telemetry_messenger_spy): - return RansomwarePayload(ransomware_payload_config, telemetry_messenger_spy) +def ransomware_payload(build_ransomware_payload, ransomware_payload_config): + return build_ransomware_payload(ransomware_payload_config) + + +@pytest.fixture +def build_ransomware_payload(telemetry_messenger_spy): + def inner(config): + return RansomwarePayload(config, telemetry_messenger_spy, shutil.copyfile) + + return inner def test_env_variables_in_target_dir_resolved_linux( - ransomware_payload_config, ransomware_target, telemetry_messenger_spy, patched_home_env + ransomware_payload_config, build_ransomware_payload, ransomware_target, patched_home_env ): path_with_env_variable = "$HOME/ransomware_target" @@ -58,7 +67,7 @@ def test_env_variables_in_target_dir_resolved_linux( ] = ransomware_payload_config["encryption"]["directories"][ "windows_target_dir" ] = path_with_env_variable - RansomwarePayload(ransomware_payload_config, telemetry_messenger_spy).run_payload() + build_ransomware_payload(ransomware_payload_config).run_payload() assert ( hash_file(ransomware_target / with_extension(ALL_ZEROS_PDF)) @@ -152,11 +161,11 @@ def test_skip_already_encrypted_file(ransomware_target, ransomware_payload): def test_encryption_skipped_if_configured_false( - ransomware_payload_config, ransomware_target, telemetry_messenger_spy + build_ransomware_payload, ransomware_payload_config, ransomware_target ): ransomware_payload_config["encryption"]["enabled"] = False - ransomware_payload = RansomwarePayload(ransomware_payload_config, telemetry_messenger_spy) + ransomware_payload = build_ransomware_payload(ransomware_payload_config) ransomware_payload.run_payload() assert hash_file(ransomware_target / ALL_ZEROS_PDF) == ALL_ZEROS_PDF_CLEARTEXT_SHA256 @@ -164,13 +173,13 @@ def test_encryption_skipped_if_configured_false( def test_encryption_skipped_if_no_directory( - ransomware_payload_config, telemetry_messenger_spy, monkeypatch + build_ransomware_payload, ransomware_payload_config, telemetry_messenger_spy ): ransomware_payload_config["encryption"]["enabled"] = True ransomware_payload_config["encryption"]["directories"]["linux_target_dir"] = "" ransomware_payload_config["encryption"]["directories"]["windows_target_dir"] = "" - ransomware_payload = RansomwarePayload(ransomware_payload_config, telemetry_messenger_spy) + ransomware_payload = build_ransomware_payload(ransomware_payload_config) ransomware_payload.run_payload() assert len(telemetry_messenger_spy.telemetries) == 0 @@ -205,17 +214,17 @@ def test_telemetry_failure(monkeypatch, ransomware_payload, telemetry_messenger_ assert "No such file or directory" in telem_1.get_data()["files"][0]["error"] -def test_readme_false(ransomware_payload_config, ransomware_target, telemetry_messenger_spy): +def test_readme_false(build_ransomware_payload, ransomware_payload_config, ransomware_target): ransomware_payload_config["other_behaviors"]["readme"] = False - ransomware_payload = RansomwarePayload(ransomware_payload_config, telemetry_messenger_spy) + ransomware_payload = build_ransomware_payload(ransomware_payload_config) ransomware_payload.run_payload() assert not Path(ransomware_target / README_DEST).exists() -def test_readme_true(ransomware_payload_config, ransomware_target, telemetry_messenger_spy): +def test_readme_true(build_ransomware_payload, ransomware_payload_config, ransomware_target): ransomware_payload_config["other_behaviors"]["readme"] = True - ransomware_payload = RansomwarePayload(ransomware_payload_config, telemetry_messenger_spy) + ransomware_payload = build_ransomware_payload(ransomware_payload_config) ransomware_payload.run_payload() assert Path(ransomware_target / README_DEST).exists() From 064525e6b9749f7da5c51497579eeadb6d3188b9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 8 Jul 2021 11:59:50 -0400 Subject: [PATCH 2/4] Agent: Don't try to create README.txt if one already exists --- .../ransomware/ransomware_payload.py | 4 ++++ .../ransomware/test_ransomware_payload.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 05c2159d2..a6637b673 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -99,6 +99,10 @@ class RansomwarePayload: def _leave_readme(self): if self._readme_enabled: readme_dest_path = self._target_dir / README_DEST + if readme_dest_path.exists(): + LOG.warning(f"{readme_dest_path} already exists, not leaving a new README.txt") + return + LOG.info(f"Leaving a ransomware README file at {readme_dest_path}") try: 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 89c9f5ff5..8a08a4595 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,6 +1,7 @@ import os import shutil from pathlib import Path, PurePosixPath +from unittest.mock import MagicMock import pytest from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import ( @@ -228,3 +229,18 @@ def test_readme_true(build_ransomware_payload, ransomware_payload_config, ransom ransomware_payload.run_payload() assert Path(ransomware_target / README_DEST).exists() + + +def test_readme_already_exists( + monkeypatch, ransomware_payload_config, telemetry_messenger_spy, ransomware_target +): + monkeypatch.setattr(ransomware_payload_module, "TARGETED_FILE_EXTENSIONS", set()), + mock_copy_file = MagicMock() + + ransomware_payload_config["other_behaviors"]["readme"] = True + Path(ransomware_target / README_DEST).touch() + RansomwarePayload( + ransomware_payload_config, telemetry_messenger_spy, mock_copy_file + ).run_payload() + + mock_copy_file.assert_not_called() From 7454ee72b229ea97b79a82ca8b4c75d4f1c5cc06 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 8 Jul 2021 12:04:12 -0400 Subject: [PATCH 3/4] Agent: Switch copy_file typehint from str to Path --- monkey/infection_monkey/ransomware/ransomware_payload.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index a6637b673..5b4a0cb8a 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -25,7 +25,7 @@ class RansomwarePayload: self, config: dict, telemetry_messenger: ITelemetryMessenger, - copy_file: Callable[[str, str], None], + copy_file: Callable[[Path, Path], None], ): LOG.debug(f"Ransomware payload configuration:\n{pformat(config)}") From 92c5c3b6826631a60a5959abc3a11a90dae877ca Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 8 Jul 2021 12:04:50 -0400 Subject: [PATCH 4/4] Agent: Extract method _copy_file() from _leave_readme() Reworks the logic in _leave_readme() to reduce indenting and improve clarity and extracts the logic to copy the file into _copy_readme_file() --- .../ransomware/ransomware_payload.py | 26 ++++++++++++------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 5b4a0cb8a..0e5cf813e 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -97,15 +97,21 @@ class RansomwarePayload: self._telemetry_messenger.send_telemetry(encryption_attempt) def _leave_readme(self): - if self._readme_enabled: - readme_dest_path = self._target_dir / README_DEST - if readme_dest_path.exists(): - LOG.warning(f"{readme_dest_path} already exists, not leaving a new README.txt") - return + if not self._readme_enabled: + return - LOG.info(f"Leaving a ransomware README file at {readme_dest_path}") + readme_dest_path = self._target_dir / README_DEST - try: - self._copy_file(README_SRC, readme_dest_path) - except Exception as ex: - LOG.warning(f"An error occurred while attempting to leave a README.txt file: {ex}") + if readme_dest_path.exists(): + LOG.warning(f"{readme_dest_path} already exists, not leaving a new README.txt") + return + + self._copy_readme_file(readme_dest_path) + + def _copy_readme_file(self, dest: Path): + LOG.info(f"Leaving a ransomware README file at {dest}") + + try: + self._copy_file(README_SRC, dest) + except Exception as ex: + LOG.warning(f"An error occurred while attempting to leave a README.txt file: {ex}")