diff --git a/.gitattributes b/.gitattributes index 807ae6822..74db1b2f8 100644 --- a/.gitattributes +++ b/.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 diff --git a/monkey/infection_monkey/i_puppet/__init__.py b/monkey/infection_monkey/i_puppet/__init__.py new file mode 100644 index 000000000..0ba1096d1 --- /dev/null +++ b/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, +) diff --git a/monkey/infection_monkey/i_puppet.py b/monkey/infection_monkey/i_puppet/i_puppet.py similarity index 89% rename from monkey/infection_monkey/i_puppet.py rename to monkey/infection_monkey/i_puppet/i_puppet.py index e25d20f53..3fa2aabd9 100644 --- a/monkey/infection_monkey/i_puppet.py +++ b/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( - self, name: str, options: Dict, interrupt: threading.Event - ) -> Tuple[None, bool, str]: + def run_payload(self, name: str, options: Dict, interrupt: threading.Event): """ 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 diff --git a/monkey/infection_monkey/puppet/plugin_type.py b/monkey/infection_monkey/i_puppet/plugin_type.py similarity index 100% rename from monkey/infection_monkey/puppet/plugin_type.py rename to monkey/infection_monkey/i_puppet/plugin_type.py diff --git a/monkey/infection_monkey/master/mock_master.py b/monkey/infection_monkey/master/mock_master.py index 0b4f9a3f6..31d4d83a7 100644 --- a/monkey/infection_monkey/master/mock_master.py +++ b/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 - path, success, error = self._puppet.run_payload("RansomwarePayload", {}, None) - self._telemetry_messenger.send_telemetry(FileEncryptionTelem(path, success, error)) + self._puppet.run_payload("RansomwarePayload", {}, None) logger.info("Finished running payloads") def terminate(self, block: bool = False) -> None: diff --git a/monkey/infection_monkey/payload/i_payload.py b/monkey/infection_monkey/payload/i_payload.py new file mode 100644 index 000000000..b63910eea --- /dev/null +++ b/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. + """ diff --git a/monkey/infection_monkey/ransomware/__init__.py b/monkey/infection_monkey/payload/ransomware/__init__.py similarity index 100% rename from monkey/infection_monkey/ransomware/__init__.py rename to monkey/infection_monkey/payload/ransomware/__init__.py diff --git a/monkey/infection_monkey/ransomware/consts.py b/monkey/infection_monkey/payload/ransomware/consts.py similarity index 100% rename from monkey/infection_monkey/ransomware/consts.py rename to monkey/infection_monkey/payload/ransomware/consts.py diff --git a/monkey/infection_monkey/ransomware/file_selectors.py b/monkey/infection_monkey/payload/ransomware/file_selectors.py similarity index 92% rename from monkey/infection_monkey/ransomware/file_selectors.py rename to monkey/infection_monkey/payload/ransomware/file_selectors.py index 33b73dd06..5707fba7d 100644 --- a/monkey/infection_monkey/ransomware/file_selectors.py +++ b/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]): diff --git a/monkey/infection_monkey/ransomware/in_place_file_encryptor.py b/monkey/infection_monkey/payload/ransomware/in_place_file_encryptor.py similarity index 100% rename from monkey/infection_monkey/ransomware/in_place_file_encryptor.py rename to monkey/infection_monkey/payload/ransomware/in_place_file_encryptor.py diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/payload/ransomware/ransomware.py similarity index 72% rename from monkey/infection_monkey/ransomware/ransomware_payload.py rename to monkey/infection_monkey/payload/ransomware/ransomware.py index a1e052970..003112cc3 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/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}") diff --git a/monkey/infection_monkey/ransomware/ransomware_payload_builder.py b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py similarity index 65% rename from monkey/infection_monkey/ransomware/ransomware_payload_builder.py rename to monkey/infection_monkey/payload/ransomware/ransomware_builder.py index 9f0d78754..4b8bbc8bb 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload_builder.py +++ b/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): - logger.debug(f"Ransomware payload configuration:\n{pformat(config)}") - ransomware_config = RansomwareConfig(config) +def build_ransomware(options: dict): + logger.debug(f"Ransomware configuration:\n{pformat(options)}") + 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( - ransomware_config, + return Ransomware( + ransomware_options, file_encryptor, file_selector, leave_readme, diff --git a/monkey/infection_monkey/ransomware/ransomware_config.py b/monkey/infection_monkey/payload/ransomware/ransomware_options.py similarity index 72% rename from monkey/infection_monkey/ransomware/ransomware_config.py rename to monkey/infection_monkey/payload/ransomware/ransomware_options.py index f8ab792da..8416f8465 100644 --- a/monkey/infection_monkey/ransomware/ransomware_config.py +++ b/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: - def __init__(self, config: dict): - self.encryption_enabled = config["encryption"]["enabled"] - self.readme_enabled = config["other_behaviors"]["readme"] +class RansomwareOptions: + def __init__(self, options: dict): + self.encryption_enabled = options["encryption"]["enabled"] + 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(): diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_payload.py b/monkey/infection_monkey/payload/ransomware/ransomware_payload.py new file mode 100644 index 000000000..d785859a2 --- /dev/null +++ b/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) diff --git a/monkey/infection_monkey/ransomware/ransomware_readme.txt b/monkey/infection_monkey/payload/ransomware/ransomware_readme.txt similarity index 100% rename from monkey/infection_monkey/ransomware/ransomware_readme.txt rename to monkey/infection_monkey/payload/ransomware/ransomware_readme.txt diff --git a/monkey/infection_monkey/ransomware/readme_dropper.py b/monkey/infection_monkey/payload/ransomware/readme_dropper.py similarity index 100% rename from monkey/infection_monkey/ransomware/readme_dropper.py rename to monkey/infection_monkey/payload/ransomware/readme_dropper.py diff --git a/monkey/infection_monkey/ransomware/targeted_file_extensions.py b/monkey/infection_monkey/payload/ransomware/targeted_file_extensions.py similarity index 100% rename from monkey/infection_monkey/ransomware/targeted_file_extensions.py rename to monkey/infection_monkey/payload/ransomware/targeted_file_extensions.py diff --git a/monkey/infection_monkey/puppet/mock_puppet.py b/monkey/infection_monkey/puppet/mock_puppet.py index 204e44ab4..182ebe55e 100644 --- a/monkey/infection_monkey/puppet/mock_puppet.py +++ b/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( - self, name: str, options: Dict, interrupt: threading.Event - ) -> Tuple[None, bool, str]: + def run_payload(self, name: str, options: Dict, interrupt: threading.Event): logger.debug(f"run_payload({name}, {options})") - return (None, True, "") def cleanup(self) -> None: print("Cleanup called!") diff --git a/monkey/infection_monkey/puppet/plugin_registry.py b/monkey/infection_monkey/puppet/plugin_registry.py new file mode 100644 index 000000000..0e98ba2ef --- /dev/null +++ b/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 diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index f932d84a4..41ed99250 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/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: - pass + def __init__(self) -> None: + 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( - self, name: str, options: Dict, interrupt: threading.Event - ) -> Tuple[None, bool, str]: - pass + def run_payload(self, name: str, options: Dict, interrupt: threading.Event): + payload = self._plugin_registry.get_plugin(name, PluginType.PAYLOAD) + payload.run(options, interrupt) def cleanup(self) -> None: pass diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/conftest.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/conftest.py similarity index 100% rename from monkey/tests/unit_tests/infection_monkey/ransomware/conftest.py rename to monkey/tests/unit_tests/infection_monkey/payload/ransomware/conftest.py diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/ransomware_target_files.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/ransomware_target_files.py similarity index 100% rename from monkey/tests/unit_tests/infection_monkey/ransomware/ransomware_target_files.py rename to monkey/tests/unit_tests/infection_monkey/payload/ransomware/ransomware_target_files.py diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_file_selectors.py similarity index 89% rename from monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py rename to monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_file_selectors.py index 42e852b95..f779b733e 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py +++ b/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.ransomware.ransomware_payload import README_SRC +from infection_monkey.payload.ransomware.file_selectors import ProductionSafeTargetFileSelector +from infection_monkey.payload.ransomware.ransomware import README_SRC TARGETED_FILE_EXTENSIONS = [".pdf", ".txt"] diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_in_place_file_encryptor.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_in_place_file_encryptor.py similarity index 92% rename from monkey/tests/unit_tests/infection_monkey/ransomware/test_in_place_file_encryptor.py rename to monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_in_place_file_encryptor.py index eb2633226..b69266db9 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_in_place_file_encryptor.py +++ b/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" diff --git a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py new file mode 100644 index 000000000..adffe6f88 --- /dev/null +++ b/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()) diff --git a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_options.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_options.py new file mode 100644 index 000000000..f2b6a8c8c --- /dev/null +++ b/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 diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_readme_dropper.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_readme_dropper.py similarity index 91% rename from monkey/tests/unit_tests/infection_monkey/ransomware/test_readme_dropper.py rename to monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_readme_dropper.py index 516e03935..8736e7c0d 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_readme_dropper.py +++ b/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" diff --git a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py new file mode 100644 index 000000000..950bc329b --- /dev/null +++ b/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() diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py deleted file mode 100644 index 141186f18..000000000 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py +++ /dev/null @@ -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 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 deleted file mode 100644 index 24eb8443d..000000000 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py +++ /dev/null @@ -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() diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_config.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_config.py index 09939b2ed..3ad02a7a6 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_config.py +++ b/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