From f0e9109f6417a5c528d7ba9a5901b819b2ff0fc2 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 8 Jul 2021 11:23:15 -0400 Subject: [PATCH] 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()