diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index e89b9ab2c..ca6976f91 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -1,7 +1,6 @@ import argparse import logging import os -import shutil import subprocess import sys import time @@ -20,6 +19,7 @@ from infection_monkey.network.HostFinger import HostFinger from infection_monkey.network.network_scanner import NetworkScanner from infection_monkey.network.tools import get_interface_to_target, is_running_on_island from infection_monkey.post_breach.post_breach_handler import PostBreach +from infection_monkey.ransomware import readme_utils from infection_monkey.ransomware.ransomware_payload import RansomwarePayload from infection_monkey.system_info import SystemInfoCollector from infection_monkey.system_singleton import SystemSingleton @@ -478,7 +478,9 @@ class InfectionMonkey(object): try: RansomwarePayload( - WormConfiguration.ransomware, batching_telemetry_messenger, shutil.copyfile + WormConfiguration.ransomware, + readme_utils.leave_readme, + batching_telemetry_messenger, ).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 f4295f866..b8f0d8805 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -24,8 +24,8 @@ class RansomwarePayload: def __init__( self, config: dict, + leave_readme: Callable[[Path, Path], None], telemetry_messenger: ITelemetryMessenger, - copy_file: Callable[[Path, Path], None], ): LOG.debug(f"Ransomware payload configuration:\n{pformat(config)}") @@ -38,7 +38,7 @@ class RansomwarePayload: self._targeted_file_extensions.discard(self._new_file_extension) self._encryptor = BitflipEncryptor(chunk_size=CHUNK_SIZE) - self._copy_file = copy_file + self._leave_readme = leave_readme self._telemetry_messenger = telemetry_messenger @staticmethod @@ -66,7 +66,7 @@ class RansomwarePayload: self._encrypt_files(file_list) if self._readme_enabled: - self._leave_readme() + self._leave_readme(README_SRC, self._target_dir / README_DEST) def _find_files(self) -> List[Path]: LOG.info(f"Collecting files in {self._target_dir}") @@ -97,21 +97,3 @@ 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(self): - - 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 - - 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}") diff --git a/monkey/infection_monkey/ransomware/readme_utils.py b/monkey/infection_monkey/ransomware/readme_utils.py new file mode 100644 index 000000000..a3037e76a --- /dev/null +++ b/monkey/infection_monkey/ransomware/readme_utils.py @@ -0,0 +1,22 @@ +import logging +import shutil +from pathlib import Path + +LOG = logging.getLogger(__name__) + + +def leave_readme(src: Path, dest: Path): + if dest.exists(): + LOG.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): + LOG.info(f"Leaving a ransomware README file at {dest}") + + try: + shutil.copyfile(src, dest) + except Exception as ex: + LOG.warning(f"An error occurred while attempting to leave a README.txt file: {ex}") diff --git a/monkey/tests/conftest.py b/monkey/tests/conftest.py index 23249016e..23cc840a3 100644 --- a/monkey/tests/conftest.py +++ b/monkey/tests/conftest.py @@ -10,4 +10,4 @@ sys.path.insert(0, MONKEY_BASE_PATH) @pytest.fixture(scope="session") def data_for_tests_dir(pytestconfig): - return os.path.join(pytestconfig.rootdir, "monkey", "tests", "data_for_tests") + return Path(os.path.join(pytestconfig.rootdir, "monkey", "tests", "data_for_tests")) diff --git a/monkey/tests/data_for_tests/test_readme.txt b/monkey/tests/data_for_tests/test_readme.txt new file mode 100644 index 000000000..8ab686eaf --- /dev/null +++ b/monkey/tests/data_for_tests/test_readme.txt @@ -0,0 +1 @@ +Hello, World! 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 2db9ecb4a..ce6b7b08e 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,5 @@ import os -import shutil -from pathlib import Path, PurePosixPath +from pathlib import PurePosixPath from unittest.mock import MagicMock import pytest @@ -24,7 +23,12 @@ from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import from tests.utils import hash_file, is_user_admin from infection_monkey.ransomware import ransomware_payload as ransomware_payload_module -from infection_monkey.ransomware.ransomware_payload import EXTENSION, README_DEST, RansomwarePayload +from infection_monkey.ransomware.ransomware_payload import ( + EXTENSION, + README_DEST, + README_SRC, + RansomwarePayload, +) def with_extension(filename): @@ -51,13 +55,18 @@ def ransomware_payload(build_ransomware_payload, ransomware_payload_config): @pytest.fixture -def build_ransomware_payload(telemetry_messenger_spy): +def build_ransomware_payload(telemetry_messenger_spy, mock_leave_readme): def inner(config): - return RansomwarePayload(config, telemetry_messenger_spy, shutil.copyfile) + return RansomwarePayload(config, mock_leave_readme, telemetry_messenger_spy) return inner +@pytest.fixture +def mock_leave_readme(): + return MagicMock() + + def test_env_variables_in_target_dir_resolved_linux( ransomware_payload_config, build_ransomware_payload, ransomware_target, patched_home_env ): @@ -215,49 +224,41 @@ 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(build_ransomware_payload, ransomware_payload_config, ransomware_target): +def test_readme_false( + build_ransomware_payload, ransomware_payload_config, mock_leave_readme, ransomware_target +): ransomware_payload_config["other_behaviors"]["readme"] = False ransomware_payload = build_ransomware_payload(ransomware_payload_config) ransomware_payload.run_payload() - assert not Path(ransomware_target / README_DEST).exists() + mock_leave_readme.assert_not_called() -def test_readme_true(build_ransomware_payload, ransomware_payload_config, ransomware_target): +def test_readme_true( + build_ransomware_payload, ransomware_payload_config, mock_leave_readme, ransomware_target +): ransomware_payload_config["other_behaviors"]["readme"] = True ransomware_payload = build_ransomware_payload(ransomware_payload_config) 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() + mock_leave_readme.assert_called_with(README_SRC, ransomware_target / README_DEST) def test_no_readme_if_no_directory( - monkeypatch, ransomware_payload_config, telemetry_messenger_spy, ransomware_target + monkeypatch, + ransomware_payload_config, + mock_leave_readme, + telemetry_messenger_spy, + ransomware_target, ): monkeypatch.setattr(ransomware_payload_module, "TARGETED_FILE_EXTENSIONS", set()), - mock_copy_file = MagicMock() ransomware_payload_config["encryption"]["directories"]["linux_target_dir"] = "" ransomware_payload_config["encryption"]["directories"]["windows_target_dir"] = "" ransomware_payload_config["other_behaviors"]["readme"] = True RansomwarePayload( - ransomware_payload_config, telemetry_messenger_spy, mock_copy_file + ransomware_payload_config, mock_leave_readme, telemetry_messenger_spy ).run_payload() - mock_copy_file.assert_not_called() + mock_leave_readme.assert_not_called() diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_readme_utils.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_readme_utils.py new file mode 100644 index 000000000..a1edf8424 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_readme_utils.py @@ -0,0 +1,32 @@ +import pytest +from tests.utils import hash_file + +from infection_monkey.ransomware.readme_utils import leave_readme + +DEST_FILE = "README.TXT" +README_HASH = "c98c24b677eff44860afea6f493bbaec5bb1c4cbb209c6fc2bbb47f66ff2ad31" +EMPTY_FILE_HASH = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + + +@pytest.fixture(scope="module") +def src_readme(data_for_tests_dir): + return data_for_tests_dir / "test_readme.txt" + + +@pytest.fixture +def dest_readme(tmp_path): + return tmp_path / DEST_FILE + + +def test_readme_already_exists(src_readme, dest_readme): + dest_readme.touch() + + leave_readme(src_readme, dest_readme) + + assert hash_file(dest_readme) == EMPTY_FILE_HASH + + +def test_leave_readme(src_readme, dest_readme): + leave_readme(src_readme, dest_readme) + + assert hash_file(dest_readme) == README_HASH