Agent: Accept a "leave_readme" Callable instead of copy_file

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}")

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}")

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}")

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"))

monkey/tests/data_for_tests/test_readme.txt

@@ -0,0 +1 @@
+Hello, World!

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()

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