Merge branch '1598-implement-run-payload' into agent-refactor

.gitattributes

@@ -1,4 +1,4 @@
 monkey/tests/data_for_tests/ransomware_targets/** -text
 monkey/tests/data_for_tests/test_readme.txt -text
 monkey/tests/data_for_tests/stable_file.txt -text
-monkey/infection_monkey/ransomware/ransomware_readme.txt -text
+monkey/infection_monkey/payload/ransomware/ransomware_readme.txt -text

monkey/infection_monkey/i_puppet/__init__.py

@@ -0,0 +1,11 @@
+from .plugin_type import PluginType
+from .i_puppet import (
+    IPuppet,
+    ExploiterResultData,
+    PingScanData,
+    PortScanData,
+    FingerprintData,
+    PortStatus,
+    PostBreachData,
+    UnknownPluginError,
+)

monkey/infection_monkey/i_puppet/i_puppet.py

@@ -2,9 +2,9 @@ import abc
 import threading
 from collections import namedtuple
 from enum import Enum
-from typing import Dict, Tuple
+from typing import Dict
 
-from infection_monkey.puppet.plugin_type import PluginType
+from . import PluginType
 
 
 class PortStatus(Enum):
@@ -27,9 +27,10 @@ PostBreachData = namedtuple("PostBreachData", ["command", "result"])
 
 class IPuppet(metaclass=abc.ABCMeta):
     @abc.abstractmethod
-    def load_plugin(self, plugin: object, plugin_type: PluginType) -> None:
+    def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None:
         """
-        Loads a plugin into the puppet.
+        Loads a plugin into the puppet
+        :param str plugin_name: The plugin class name
         :param object plugin: The plugin object to load
         :param PluginType plugin_type: The type of plugin being loaded
         """
@@ -107,13 +108,13 @@ class IPuppet(metaclass=abc.ABCMeta):
         """
 
     @abc.abstractmethod
-    def run_payload(
+    def run_payload(self, name: str, options: Dict, interrupt: threading.Event):
-        self, name: str, options: Dict, interrupt: threading.Event
-    ) -> Tuple[None, bool, str]:
         """
         Runs a payload
         :param str name: The name of the payload to run
         :param Dict options: A dictionary containing options that modify the behavior of the payload
+        :param threading.Event interrupt: A threading.Event object that signals the payload to stop
+                                          executing and clean itself up.
         """
 
     @abc.abstractmethod

monkey/infection_monkey/master/mock_master.py

@@ -4,7 +4,6 @@ from infection_monkey.i_master import IMaster
 from infection_monkey.i_puppet import IPuppet, PortStatus
 from infection_monkey.model.host import VictimHost
 from infection_monkey.telemetry.exploit_telem import ExploitTelem
-from infection_monkey.telemetry.file_encryption_telem import FileEncryptionTelem
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
 from infection_monkey.telemetry.scan_telem import ScanTelem
@@ -119,9 +118,7 @@ class MockMaster(IMaster):
 
     def _run_payload(self):
         logger.info("Running payloads")
-        # TODO: modify what FileEncryptionTelem gets
+        self._puppet.run_payload("RansomwarePayload", {}, None)
-        path, success, error = self._puppet.run_payload("RansomwarePayload", {}, None)
-        self._telemetry_messenger.send_telemetry(FileEncryptionTelem(path, success, error))
         logger.info("Finished running payloads")
 
     def terminate(self, block: bool = False) -> None:

monkey/infection_monkey/payload/i_payload.py

@@ -0,0 +1,14 @@
+import abc
+import threading
+from typing import Dict
+
+
+class IPayload(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def run(self, options: Dict, interrupt: threading.Event):
+        """
+        Runs the payload
+        :param Dict options: A dictionary containing options that modify the behavior of the payload
+        :param threading.Event interrupt: A threading.Event object that signals the payload to stop
+                                          executing and clean itself up.
+        """

monkey/infection_monkey/payload/ransomware/file_selectors.py

@@ -2,7 +2,6 @@ from pathlib import Path
 from typing import List, Set
 
 from common.utils.file_utils import get_file_sha256_hash
-from infection_monkey.ransomware.consts import README_FILE_NAME, README_SHA256_HASH
 from infection_monkey.utils.dir_utils import (
     file_extension_filter,
     filter_files,
@@ -11,6 +10,8 @@ from infection_monkey.utils.dir_utils import (
     is_not_symlink_filter,
 )
 
+from .consts import README_FILE_NAME, README_SHA256_HASH
+
 
 class ProductionSafeTargetFileSelector:
     def __init__(self, targeted_file_extensions: Set[str]):

monkey/infection_monkey/payload/ransomware/ransomware.py

@@ -1,19 +1,21 @@
 import logging
+import threading
 from pathlib import Path
 from typing import Callable, List
 
-from infection_monkey.ransomware.consts import README_FILE_NAME, README_SRC
-from infection_monkey.ransomware.ransomware_config import RansomwareConfig
 from infection_monkey.telemetry.file_encryption_telem import FileEncryptionTelem
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 
+from .consts import README_FILE_NAME, README_SRC
+from .ransomware_options import RansomwareOptions
+
 logger = logging.getLogger(__name__)
 
 
-class RansomwarePayload:
+class Ransomware:
     def __init__(
         self,
-        config: RansomwareConfig,
+        config: RansomwareOptions,
         encrypt_file: Callable[[Path], None],
         select_files: Callable[[Path], List[Path]],
         leave_readme: Callable[[Path, Path], None],
@@ -31,7 +33,7 @@ class RansomwarePayload:
             self._target_directory / README_FILE_NAME if self._target_directory else None
         )
 
-    def run_payload(self):
+    def run(self, interrupt: threading.Event):
         if not self._target_directory:
             return
 
@@ -39,19 +41,25 @@ class RansomwarePayload:
 
         if self._config.encryption_enabled:
             file_list = self._find_files()
-            self._encrypt_files(file_list)
+            self._encrypt_files(file_list, interrupt)
 
         if self._config.readme_enabled:
-            self._leave_readme_in_target_directory()
+            self._leave_readme_in_target_directory(interrupt)
 
     def _find_files(self) -> List[Path]:
         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]):
+    def _encrypt_files(self, file_list: List[Path], interrupt: threading.Event):
         logger.info(f"Encrypting files in {self._target_directory}")
 
         for filepath in file_list:
+            if interrupt.is_set():
+                logger.debug(
+                    "Received a stop signal, skipping remaining files for encryption of "
+                    "ransomware payload"
+                )
+                return
             try:
                 logger.debug(f"Encrypting {filepath}")
                 self._encrypt_file(filepath)
@@ -64,8 +72,12 @@ class RansomwarePayload:
         encryption_attempt = FileEncryptionTelem(str(filepath), success, error)
         self._telemetry_messenger.send_telemetry(encryption_attempt)
 
-    def _leave_readme_in_target_directory(self):
+    def _leave_readme_in_target_directory(self, interrupt: threading.Event):
         try:
+            if interrupt.is_set():
+                logger.debug("Received a stop signal, skipping leave readme")
+                return
+
             self._leave_readme(README_SRC, self._readme_file_path)
         except Exception as ex:
             logger.warning(f"An error occurred while attempting to leave a README.txt file: {ex}")

monkey/infection_monkey/payload/ransomware/ransomware_builder.py

@@ -1,12 +1,6 @@
 import logging
 from pprint import pformat
 
-from infection_monkey.ransomware import readme_dropper
-from infection_monkey.ransomware.file_selectors import ProductionSafeTargetFileSelector
-from infection_monkey.ransomware.in_place_file_encryptor import InPlaceFileEncryptor
-from infection_monkey.ransomware.ransomware_config import RansomwareConfig
-from infection_monkey.ransomware.ransomware_payload import RansomwarePayload
-from infection_monkey.ransomware.targeted_file_extensions import TARGETED_FILE_EXTENSIONS
 from infection_monkey.telemetry.messengers.batching_telemetry_messenger import (
     BatchingTelemetryMessenger,
 )
@@ -15,23 +9,30 @@ from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter im
 )
 from infection_monkey.utils.bit_manipulators import flip_bits
 
+from . import readme_dropper
+from .file_selectors import ProductionSafeTargetFileSelector
+from .in_place_file_encryptor import InPlaceFileEncryptor
+from .ransomware import Ransomware
+from .ransomware_options import RansomwareOptions
+from .targeted_file_extensions import TARGETED_FILE_EXTENSIONS
+
 EXTENSION = ".m0nk3y"
 CHUNK_SIZE = 4096 * 24
 
 logger = logging.getLogger(__name__)
 
 
-def build_ransomware_payload(config: dict):
+def build_ransomware(options: dict):
-    logger.debug(f"Ransomware payload configuration:\n{pformat(config)}")
+    logger.debug(f"Ransomware configuration:\n{pformat(options)}")
-    ransomware_config = RansomwareConfig(config)
+    ransomware_options = RansomwareOptions(options)
 
     file_encryptor = _build_file_encryptor()
     file_selector = _build_file_selector()
     leave_readme = _build_leave_readme()
     telemetry_messenger = _build_telemetry_messenger()
 
-    return RansomwarePayload(
+    return Ransomware(
-        ransomware_config,
+        ransomware_options,
         file_encryptor,
         file_selector,
         leave_readme,

monkey/infection_monkey/payload/ransomware/ransomware_options.py

@@ -6,13 +6,13 @@ from infection_monkey.utils.environment import is_windows_os
 logger = logging.getLogger(__name__)
 
 
-class RansomwareConfig:
+class RansomwareOptions:
-    def __init__(self, config: dict):
+    def __init__(self, options: dict):
-        self.encryption_enabled = config["encryption"]["enabled"]
+        self.encryption_enabled = options["encryption"]["enabled"]
-        self.readme_enabled = config["other_behaviors"]["readme"]
+        self.readme_enabled = options["other_behaviors"]["readme"]
 
         self.target_directory = None
-        self._set_target_directory(config["encryption"]["directories"])
+        self._set_target_directory(options["encryption"]["directories"])
 
     def _set_target_directory(self, os_target_directories: dict):
         if is_windows_os():

monkey/infection_monkey/payload/ransomware/ransomware_payload.py

@@ -0,0 +1,12 @@
+import threading
+from typing import Dict
+
+from infection_monkey.payload.i_payload import IPayload
+
+from . import ransomware_builder
+
+
+class RansomwarePayload(IPayload):
+    def run(self, options: Dict, interrupt: threading.Event):
+        ransomware = ransomware_builder.build_ransomware(options)
+        ransomware.run(interrupt)

monkey/infection_monkey/puppet/mock_puppet.py

@@ -1,17 +1,17 @@
 import logging
 import threading
-from typing import Dict, Tuple
+from typing import Dict
 
 from infection_monkey.i_puppet import (
     ExploiterResultData,
     FingerprintData,
     IPuppet,
     PingScanData,
+    PluginType,
     PortScanData,
     PortStatus,
     PostBreachData,
 )
-from infection_monkey.puppet.plugin_type import PluginType
 
 DOT_1 = "10.0.0.1"
 DOT_2 = "10.0.0.2"
@@ -299,11 +299,8 @@ class MockPuppet(IPuppet):
         except KeyError:
             return ExploiterResultData(False, {}, [], f"{name} failed for host {host}")
 
-    def run_payload(
+    def run_payload(self, name: str, options: Dict, interrupt: threading.Event):
-        self, name: str, options: Dict, interrupt: threading.Event
-    ) -> Tuple[None, bool, str]:
         logger.debug(f"run_payload({name}, {options})")
-        return (None, True, "")
 
     def cleanup(self) -> None:
         print("Cleanup called!")

monkey/infection_monkey/puppet/plugin_registry.py

@@ -0,0 +1,41 @@
+import logging
+from typing import Optional
+
+from infection_monkey.i_puppet import PluginType, UnknownPluginError
+
+logger = logging.getLogger()
+
+
+class PluginRegistry:
+    def __init__(self):
+        """
+        `self._registry` looks like -
+            {
+                PluginType.EXPLOITER: {
+                    "ZerologonExploiter": ZerologonExploiter,
+                    "SMBExploiter": SMBExploiter
+                },
+                PluginType.PBA: {
+                    "CommunicateAsBackdoorUser": CommunicateAsBackdoorUser
+                }
+            }
+        """
+        self._registry = {}
+
+    def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None:
+        self._registry.setdefault(plugin_type, {})
+        self._registry[plugin_type][plugin_name] = plugin
+
+        logger.debug(f"Plugin '{plugin_name}' loaded")
+
+    def get_plugin(self, plugin_name: str, plugin_type: PluginType) -> Optional[object]:
+        try:
+            plugin = self._registry[plugin_type][plugin_name]
+        except KeyError:
+            raise UnknownPluginError(
+                f"Unknown plugin '{plugin_name}' of type '{plugin_type.value}'"
+            )
+
+        logger.debug(f"Plugin '{plugin_name}' found")
+
+        return plugin

monkey/infection_monkey/puppet/puppet.py

@@ -1,35 +1,42 @@
 import logging
 import threading
-from typing import Dict, Tuple
+from typing import Dict
 
 from infection_monkey.i_puppet import (
     ExploiterResultData,
     FingerprintData,
     IPuppet,
     PingScanData,
+    PluginType,
     PortScanData,
     PostBreachData,
 )
-from infection_monkey.puppet.plugin_type import PluginType
+
+from .mock_puppet import MockPuppet
+from .plugin_registry import PluginRegistry
 
 logger = logging.getLogger()
 
 
 class Puppet(IPuppet):
-    def load_plugin(self, plugin: object, plugin_type: PluginType) -> None:
+    def __init__(self) -> None:
-        pass
+        self._mock_puppet = MockPuppet()
+        self._plugin_registry = PluginRegistry()
+
+    def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None:
+        self._plugin_registry.load_plugin(plugin_name, plugin, plugin_type)
 
     def run_sys_info_collector(self, name: str) -> Dict:
-        pass
+        return self._mock_puppet.run_sys_info_collector(name)
 
     def run_pba(self, name: str, options: Dict) -> PostBreachData:
-        pass
+        return self._mock_puppet.run_pba(name, options)
 
     def ping(self, host: str, timeout: float = 1) -> PingScanData:
-        pass
+        return self._mock_puppet.ping(host, timeout)
 
     def scan_tcp_port(self, host: str, port: int, timeout: float = 3) -> PortScanData:
-        pass
+        return self._mock_puppet.scan_tcp_port(host, port, timeout)
 
     def fingerprint(
         self,
@@ -38,17 +45,16 @@ class Puppet(IPuppet):
         ping_scan_data: PingScanData,
         port_scan_data: Dict[int, PortScanData],
     ) -> FingerprintData:
-        pass
+        return self._mock_puppet.fingerprint(name, host, ping_scan_data, port_scan_data)
 
     def exploit_host(
         self, name: str, host: str, options: Dict, interrupt: threading.Event
     ) -> ExploiterResultData:
-        pass
+        return self._mock_puppet.exploit_host(name, host, options, interrupt)
 
-    def run_payload(
+    def run_payload(self, name: str, options: Dict, interrupt: threading.Event):
-        self, name: str, options: Dict, interrupt: threading.Event
+        payload = self._plugin_registry.get_plugin(name, PluginType.PAYLOAD)
-    ) -> Tuple[None, bool, str]:
+        payload.run(options, interrupt)
-        pass
 
     def cleanup(self) -> None:
         pass

monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_file_selectors.py

@@ -2,7 +2,7 @@ import os
 import shutil
 
 import pytest
-from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import (
+from tests.unit_tests.infection_monkey.payload.ransomware.ransomware_target_files import (
     ALL_ZEROS_PDF,
     HELLO_TXT,
     SHORTCUT_LNK,
@@ -12,8 +12,8 @@ from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import
 )
 from tests.utils import is_user_admin
 
-from infection_monkey.ransomware.file_selectors import ProductionSafeTargetFileSelector
+from infection_monkey.payload.ransomware.file_selectors import ProductionSafeTargetFileSelector
-from infection_monkey.ransomware.ransomware_payload import README_SRC
+from infection_monkey.payload.ransomware.ransomware import README_SRC
 
 TARGETED_FILE_EXTENSIONS = [".pdf", ".txt"]
 

monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_in_place_file_encryptor.py

@@ -1,7 +1,7 @@
 import os
 
 import pytest
-from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import (
+from tests.unit_tests.infection_monkey.payload.ransomware.ransomware_target_files import (
     ALL_ZEROS_PDF,
     ALL_ZEROS_PDF_CLEARTEXT_SHA256,
     ALL_ZEROS_PDF_ENCRYPTED_SHA256,
@@ -11,7 +11,7 @@ from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import
 )
 
 from common.utils.file_utils import get_file_sha256_hash
-from infection_monkey.ransomware.in_place_file_encryptor import InPlaceFileEncryptor
+from infection_monkey.payload.ransomware.in_place_file_encryptor import InPlaceFileEncryptor
 from infection_monkey.utils.bit_manipulators import flip_bits
 
 EXTENSION = ".m0nk3y"

monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py

@@ -0,0 +1,213 @@
+import threading
+from pathlib import PurePosixPath
+from unittest.mock import MagicMock
+
+import pytest
+from tests.unit_tests.infection_monkey.payload.ransomware.ransomware_target_files import (
+    ALL_ZEROS_PDF,
+    HELLO_TXT,
+    TEST_KEYBOARD_TXT,
+)
+
+from infection_monkey.payload.ransomware.consts import README_FILE_NAME, README_SRC
+from infection_monkey.payload.ransomware.ransomware import Ransomware
+from infection_monkey.payload.ransomware.ransomware_options import RansomwareOptions
+
+
+@pytest.fixture
+def ransomware(build_ransomware, ransomware_options):
+    return build_ransomware(ransomware_options)
+
+
+@pytest.fixture
+def build_ransomware(
+    mock_file_encryptor, mock_file_selector, mock_leave_readme, telemetry_messenger_spy
+):
+    def inner(
+        config,
+        file_encryptor=mock_file_encryptor,
+        file_selector=mock_file_selector,
+        leave_readme=mock_leave_readme,
+    ):
+        return Ransomware(
+            config,
+            file_encryptor,
+            file_selector,
+            leave_readme,
+            telemetry_messenger_spy,
+        )
+
+    return inner
+
+
+@pytest.fixture
+def ransomware_options(ransomware_test_data):
+    class RansomwareOptionsStub(RansomwareOptions):
+        def __init__(self, encryption_enabled, readme_enabled, target_directory):
+            self.encryption_enabled = encryption_enabled
+            self.readme_enabled = readme_enabled
+            self.target_directory = target_directory
+
+    return RansomwareOptionsStub(True, False, ransomware_test_data)
+
+
+@pytest.fixture
+def mock_file_encryptor():
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_file_selector(ransomware_test_data):
+    selected_files = [
+        ransomware_test_data / ALL_ZEROS_PDF,
+        ransomware_test_data / TEST_KEYBOARD_TXT,
+    ]
+    return MagicMock(return_value=selected_files)
+
+
+@pytest.fixture
+def mock_leave_readme():
+    return MagicMock()
+
+
+@pytest.fixture
+def interrupt():
+    return threading.Event()
+
+
+def test_files_selected_from_target_dir(
+    ransomware,
+    ransomware_options,
+    mock_file_selector,
+):
+    ransomware.run(threading.Event())
+    mock_file_selector.assert_called_with(ransomware_options.target_directory)
+
+
+def test_all_selected_files_encrypted(ransomware_test_data, ransomware, mock_file_encryptor):
+    ransomware.run(threading.Event())
+
+    assert mock_file_encryptor.call_count == 2
+    mock_file_encryptor.assert_any_call(ransomware_test_data / ALL_ZEROS_PDF)
+    mock_file_encryptor.assert_any_call(ransomware_test_data / TEST_KEYBOARD_TXT)
+
+
+def test_interrupt_while_encrypting(
+    ransomware_test_data, interrupt, ransomware_options, build_ransomware
+):
+    selected_files = [
+        ransomware_test_data / ALL_ZEROS_PDF,
+        ransomware_test_data / HELLO_TXT,
+        ransomware_test_data / TEST_KEYBOARD_TXT,
+    ]
+    mfs = MagicMock(return_value=selected_files)
+
+    def _callback(file_path, *_):
+        # Block all threads here until 2 threads reach this barrier, then set stop
+        # and test that neither thread continues to scan.
+        if file_path.name == HELLO_TXT:
+            interrupt.set()
+
+    mfe = MagicMock(side_effect=_callback)
+
+    build_ransomware(ransomware_options, mfe, mfs).run(interrupt)
+
+    assert mfe.call_count == 2
+    mfe.assert_any_call(ransomware_test_data / ALL_ZEROS_PDF)
+    mfe.assert_any_call(ransomware_test_data / HELLO_TXT)
+
+
+def test_no_readme_after_interrupt(ransomware, interrupt, mock_leave_readme):
+    interrupt.set()
+    ransomware.run(interrupt)
+
+    mock_leave_readme.assert_not_called()
+
+
+def test_encryption_skipped_if_configured_false(
+    build_ransomware, ransomware_options, mock_file_encryptor
+):
+    ransomware_options.encryption_enabled = False
+
+    ransomware = build_ransomware(ransomware_options)
+    ransomware.run(threading.Event())
+
+    assert mock_file_encryptor.call_count == 0
+
+
+def test_encryption_skipped_if_no_directory(
+    build_ransomware, ransomware_options, mock_file_encryptor
+):
+    ransomware_options.encryption_enabled = True
+    ransomware_options.target_directory = None
+
+    ransomware = build_ransomware(ransomware_options)
+    ransomware.run(threading.Event())
+
+    assert mock_file_encryptor.call_count == 0
+
+
+def test_telemetry_success(ransomware, telemetry_messenger_spy):
+    ransomware.run(threading.Event())
+
+    assert len(telemetry_messenger_spy.telemetries) == 2
+    telem_1 = telemetry_messenger_spy.telemetries[0]
+    telem_2 = telemetry_messenger_spy.telemetries[1]
+
+    assert ALL_ZEROS_PDF in telem_1.get_data()["files"][0]["path"]
+    assert telem_1.get_data()["files"][0]["success"]
+    assert telem_1.get_data()["files"][0]["error"] == ""
+    assert TEST_KEYBOARD_TXT in telem_2.get_data()["files"][0]["path"]
+    assert telem_2.get_data()["files"][0]["success"]
+    assert telem_2.get_data()["files"][0]["error"] == ""
+
+
+def test_telemetry_failure(build_ransomware, ransomware_options, telemetry_messenger_spy):
+    file_not_exists = "/file/not/exist"
+    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 = build_ransomware(config=ransomware_options, file_encryptor=mfe, file_selector=mfs)
+
+    ransomware.run(threading.Event())
+    telem = telemetry_messenger_spy.telemetries[0]
+
+    assert file_not_exists in telem.get_data()["files"][0]["path"]
+    assert not telem.get_data()["files"][0]["success"]
+    assert "No such file or directory" in telem.get_data()["files"][0]["error"]
+
+
+def test_readme_false(build_ransomware, ransomware_options, mock_leave_readme):
+    ransomware_options.readme_enabled = False
+    ransomware = build_ransomware(ransomware_options)
+
+    ransomware.run(threading.Event())
+    mock_leave_readme.assert_not_called()
+
+
+def test_readme_true(build_ransomware, ransomware_options, mock_leave_readme, ransomware_test_data):
+    ransomware_options.readme_enabled = True
+    ransomware = build_ransomware(ransomware_options)
+
+    ransomware.run(threading.Event())
+    mock_leave_readme.assert_called_with(README_SRC, ransomware_test_data / README_FILE_NAME)
+
+
+def test_no_readme_if_no_directory(build_ransomware, ransomware_options, mock_leave_readme):
+    ransomware_options.target_directory = None
+    ransomware_options.readme_enabled = True
+
+    ransomware = build_ransomware(ransomware_options)
+
+    ransomware.run(threading.Event())
+    mock_leave_readme.assert_not_called()
+
+
+def test_leave_readme_exceptions_handled(build_ransomware, ransomware_options):
+    leave_readme = MagicMock(side_effect=Exception("Test exception when leaving README"))
+    ransomware_options.readme_enabled = True
+    ransomware = build_ransomware(config=ransomware_options, leave_readme=leave_readme)
+
+    # Test will fail if exception is raised and not handled
+    ransomware.run(threading.Event())

monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_options.py

@@ -0,0 +1,73 @@
+from pathlib import Path
+
+import pytest
+from tests.utils import raise_
+
+from common.utils.file_utils import InvalidPath
+from infection_monkey.payload.ransomware import ransomware_options
+from infection_monkey.payload.ransomware.ransomware_options import RansomwareOptions
+
+LINUX_DIR = "/tmp/test"
+WINDOWS_DIR = "C:\\tmp\\test"
+
+
+@pytest.fixture
+def options_from_island():
+    return {
+        "encryption": {
+            "enabled": None,
+            "directories": {
+                "linux_target_dir": LINUX_DIR,
+                "windows_target_dir": WINDOWS_DIR,
+            },
+        },
+        "other_behaviors": {"readme": None},
+    }
+
+
+@pytest.mark.parametrize("enabled", [True, False])
+def test_encryption_enabled(enabled, options_from_island):
+    options_from_island["encryption"]["enabled"] = enabled
+    options = RansomwareOptions(options_from_island)
+
+    assert options.encryption_enabled == enabled
+
+
+@pytest.mark.parametrize("enabled", [True, False])
+def test_readme_enabled(enabled, options_from_island):
+    options_from_island["other_behaviors"]["readme"] = enabled
+    options = RansomwareOptions(options_from_island)
+
+    assert options.readme_enabled == enabled
+
+
+def test_linux_target_dir(monkeypatch, options_from_island):
+    monkeypatch.setattr(ransomware_options, "is_windows_os", lambda: False)
+
+    options = RansomwareOptions(options_from_island)
+    assert options.target_directory == Path(LINUX_DIR)
+
+
+def test_windows_target_dir(monkeypatch, options_from_island):
+    monkeypatch.setattr(ransomware_options, "is_windows_os", lambda: True)
+
+    options = RansomwareOptions(options_from_island)
+    assert options.target_directory == Path(WINDOWS_DIR)
+
+
+def test_env_variables_in_target_dir_resolved(options_from_island, patched_home_env, tmp_path):
+    path_with_env_variable = "$HOME/ransomware_target"
+
+    options_from_island["encryption"]["directories"]["linux_target_dir"] = options_from_island[
+        "encryption"
+    ]["directories"]["windows_target_dir"] = path_with_env_variable
+
+    options = RansomwareOptions(options_from_island)
+    assert options.target_directory == patched_home_env / "ransomware_target"
+
+
+def test_target_dir_is_none(monkeypatch, options_from_island):
+    monkeypatch.setattr(ransomware_options, "expand_path", lambda _: raise_(InvalidPath("invalid")))
+
+    options = RansomwareOptions(options_from_island)
+    assert options.target_directory is None

monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_readme_dropper.py

@@ -1,7 +1,7 @@
 import pytest
 
 from common.utils.file_utils import get_file_sha256_hash
-from infection_monkey.ransomware.readme_dropper import leave_readme
+from infection_monkey.payload.ransomware.readme_dropper import leave_readme
 
 DEST_FILE = "README.TXT"
 README_HASH = "c98c24b677eff44860afea6f493bbaec5bb1c4cbb209c6fc2bbb47f66ff2ad31"

monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py

@@ -0,0 +1,43 @@
+import threading
+from unittest.mock import MagicMock
+
+from infection_monkey.i_puppet import PluginType
+from infection_monkey.puppet.puppet import Puppet
+
+
+def test_puppet_run_payload_success(monkeypatch):
+    p = Puppet()
+
+    payload = MagicMock()
+    payload_name = "PayloadOne"
+
+    p.load_plugin(payload_name, payload, PluginType.PAYLOAD)
+    p.run_payload(payload_name, {}, threading.Event())
+
+    payload.run.assert_called_once()
+
+
+def test_puppet_run_multiple_payloads(monkeypatch):
+    p = Puppet()
+
+    payload_1 = MagicMock()
+    payload1_name = "PayloadOne"
+
+    payload_2 = MagicMock()
+    payload2_name = "PayloadTwo"
+
+    payload_3 = MagicMock()
+    payload3_name = "PayloadThree"
+
+    p.load_plugin(payload1_name, payload_1, PluginType.PAYLOAD)
+    p.load_plugin(payload2_name, payload_2, PluginType.PAYLOAD)
+    p.load_plugin(payload3_name, payload_3, PluginType.PAYLOAD)
+
+    p.run_payload(payload1_name, {}, threading.Event())
+    payload_1.run.assert_called_once()
+
+    p.run_payload(payload2_name, {}, threading.Event())
+    payload_2.run.assert_called_once()
+
+    p.run_payload(payload3_name, {}, threading.Event())
+    payload_3.run.assert_called_once()

monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py

@@ -1,73 +0,0 @@
-from pathlib import Path
-
-import pytest
-from tests.utils import raise_
-
-from common.utils.file_utils import InvalidPath
-from infection_monkey.ransomware import ransomware_config
-from infection_monkey.ransomware.ransomware_config import RansomwareConfig
-
-LINUX_DIR = "/tmp/test"
-WINDOWS_DIR = "C:\\tmp\\test"
-
-
-@pytest.fixture
-def config_from_island():
-    return {
-        "encryption": {
-            "enabled": None,
-            "directories": {
-                "linux_target_dir": LINUX_DIR,
-                "windows_target_dir": WINDOWS_DIR,
-            },
-        },
-        "other_behaviors": {"readme": None},
-    }
-
-
-@pytest.mark.parametrize("enabled", [True, False])
-def test_encryption_enabled(enabled, config_from_island):
-    config_from_island["encryption"]["enabled"] = enabled
-    config = RansomwareConfig(config_from_island)
-
-    assert config.encryption_enabled == enabled
-
-
-@pytest.mark.parametrize("enabled", [True, False])
-def test_readme_enabled(enabled, config_from_island):
-    config_from_island["other_behaviors"]["readme"] = enabled
-    config = RansomwareConfig(config_from_island)
-
-    assert config.readme_enabled == enabled
-
-
-def test_linux_target_dir(monkeypatch, config_from_island):
-    monkeypatch.setattr(ransomware_config, "is_windows_os", lambda: False)
-
-    config = RansomwareConfig(config_from_island)
-    assert config.target_directory == Path(LINUX_DIR)
-
-
-def test_windows_target_dir(monkeypatch, config_from_island):
-    monkeypatch.setattr(ransomware_config, "is_windows_os", lambda: True)
-
-    config = RansomwareConfig(config_from_island)
-    assert config.target_directory == Path(WINDOWS_DIR)
-
-
-def test_env_variables_in_target_dir_resolved(config_from_island, patched_home_env, tmp_path):
-    path_with_env_variable = "$HOME/ransomware_target"
-
-    config_from_island["encryption"]["directories"]["linux_target_dir"] = config_from_island[
-        "encryption"
-    ]["directories"]["windows_target_dir"] = path_with_env_variable
-
-    config = RansomwareConfig(config_from_island)
-    assert config.target_directory == patched_home_env / "ransomware_target"
-
-
-def test_target_dir_is_none(monkeypatch, config_from_island):
-    monkeypatch.setattr(ransomware_config, "expand_path", lambda _: raise_(InvalidPath("invalid")))
-
-    config = RansomwareConfig(config_from_island)
-    assert config.target_directory is None

monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py

@@ -1,186 +0,0 @@
-from pathlib import PurePosixPath
-from unittest.mock import MagicMock
-
-import pytest
-from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import (
-    ALL_ZEROS_PDF,
-    TEST_KEYBOARD_TXT,
-)
-
-from infection_monkey.ransomware.consts import README_FILE_NAME, README_SRC
-from infection_monkey.ransomware.ransomware_config import RansomwareConfig
-from infection_monkey.ransomware.ransomware_payload import RansomwarePayload
-
-
-@pytest.fixture
-def ransomware_payload(build_ransomware_payload, ransomware_payload_config):
-    return build_ransomware_payload(ransomware_payload_config)
-
-
-@pytest.fixture
-def build_ransomware_payload(
-    mock_file_encryptor, mock_file_selector, mock_leave_readme, telemetry_messenger_spy
-):
-    def inner(
-        config,
-        file_encryptor=mock_file_encryptor,
-        file_selector=mock_file_selector,
-        leave_readme=mock_leave_readme,
-    ):
-        return RansomwarePayload(
-            config,
-            file_encryptor,
-            file_selector,
-            leave_readme,
-            telemetry_messenger_spy,
-        )
-
-    return inner
-
-
-@pytest.fixture
-def ransomware_payload_config(ransomware_test_data):
-    class RansomwareConfigStub(RansomwareConfig):
-        def __init__(self, encryption_enabled, readme_enabled, target_directory):
-            self.encryption_enabled = encryption_enabled
-            self.readme_enabled = readme_enabled
-            self.target_directory = target_directory
-
-    return RansomwareConfigStub(True, False, ransomware_test_data)
-
-
-@pytest.fixture
-def mock_file_encryptor():
-    return MagicMock()
-
-
-@pytest.fixture
-def mock_file_selector(ransomware_test_data):
-    selected_files = [
-        ransomware_test_data / ALL_ZEROS_PDF,
-        ransomware_test_data / TEST_KEYBOARD_TXT,
-    ]
-    return MagicMock(return_value=selected_files)
-
-
-@pytest.fixture
-def mock_leave_readme():
-    return MagicMock()
-
-
-def test_files_selected_from_target_dir(
-    ransomware_payload,
-    ransomware_payload_config,
-    mock_file_selector,
-):
-    ransomware_payload.run_payload()
-    mock_file_selector.assert_called_with(ransomware_payload_config.target_directory)
-
-
-def test_all_selected_files_encrypted(
-    ransomware_test_data, ransomware_payload, mock_file_encryptor
-):
-    ransomware_payload.run_payload()
-
-    assert mock_file_encryptor.call_count == 2
-    mock_file_encryptor.assert_any_call(ransomware_test_data / ALL_ZEROS_PDF)
-    mock_file_encryptor.assert_any_call(ransomware_test_data / TEST_KEYBOARD_TXT)
-
-
-def test_encryption_skipped_if_configured_false(
-    build_ransomware_payload, ransomware_payload_config, mock_file_encryptor
-):
-    ransomware_payload_config.encryption_enabled = False
-
-    ransomware_payload = build_ransomware_payload(ransomware_payload_config)
-    ransomware_payload.run_payload()
-
-    assert mock_file_encryptor.call_count == 0
-
-
-def test_encryption_skipped_if_no_directory(
-    build_ransomware_payload, ransomware_payload_config, mock_file_encryptor
-):
-    ransomware_payload_config.encryption_enabled = True
-    ransomware_payload_config.target_directory = None
-
-    ransomware_payload = build_ransomware_payload(ransomware_payload_config)
-    ransomware_payload.run_payload()
-
-    assert mock_file_encryptor.call_count == 0
-
-
-def test_telemetry_success(ransomware_payload, telemetry_messenger_spy):
-    ransomware_payload.run_payload()
-
-    assert len(telemetry_messenger_spy.telemetries) == 2
-    telem_1 = telemetry_messenger_spy.telemetries[0]
-    telem_2 = telemetry_messenger_spy.telemetries[1]
-
-    assert ALL_ZEROS_PDF in telem_1.get_data()["files"][0]["path"]
-    assert telem_1.get_data()["files"][0]["success"]
-    assert telem_1.get_data()["files"][0]["error"] == ""
-    assert TEST_KEYBOARD_TXT in telem_2.get_data()["files"][0]["path"]
-    assert telem_2.get_data()["files"][0]["success"]
-    assert telem_2.get_data()["files"][0]["error"] == ""
-
-
-def test_telemetry_failure(
-    build_ransomware_payload, ransomware_payload_config, telemetry_messenger_spy
-):
-    file_not_exists = "/file/not/exist"
-    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()
-    telem = telemetry_messenger_spy.telemetries[0]
-
-    assert file_not_exists in telem.get_data()["files"][0]["path"]
-    assert not telem.get_data()["files"][0]["success"]
-    assert "No such file or directory" in telem.get_data()["files"][0]["error"]
-
-
-def test_readme_false(build_ransomware_payload, ransomware_payload_config, mock_leave_readme):
-    ransomware_payload_config.readme_enabled = False
-    ransomware_payload = build_ransomware_payload(ransomware_payload_config)
-
-    ransomware_payload.run_payload()
-    mock_leave_readme.assert_not_called()
-
-
-def test_readme_true(
-    build_ransomware_payload, ransomware_payload_config, mock_leave_readme, ransomware_test_data
-):
-    ransomware_payload_config.readme_enabled = True
-    ransomware_payload = build_ransomware_payload(ransomware_payload_config)
-
-    ransomware_payload.run_payload()
-    mock_leave_readme.assert_called_with(README_SRC, ransomware_test_data / README_FILE_NAME)
-
-
-def test_no_readme_if_no_directory(
-    build_ransomware_payload, ransomware_payload_config, mock_leave_readme
-):
-    ransomware_payload_config.target_directory = None
-    ransomware_payload_config.readme_enabled = True
-
-    ransomware_payload = build_ransomware_payload(ransomware_payload_config)
-
-    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()

monkey/tests/unit_tests/monkey_island/cc/services/test_config.py

@@ -38,7 +38,7 @@ def test_format_config_for_agent__credentials_removed(flat_monkey_config):
 
 
 def test_format_config_for_agent__ransomware_payload(flat_monkey_config):
-    expected_ransomware_config = {
+    expected_ransomware_options = {
         "ransomware": {
             "encryption": {
                 "enabled": True,
@@ -54,7 +54,7 @@ def test_format_config_for_agent__ransomware_payload(flat_monkey_config):
     ConfigService.format_flat_config_for_agent(flat_monkey_config)
 
     assert "payloads" in flat_monkey_config
-    assert flat_monkey_config["payloads"] == expected_ransomware_config
+    assert flat_monkey_config["payloads"] == expected_ransomware_options
 
     assert "ransomware" not in flat_monkey_config