From 3adb1d5b071b5f7ac86e3b39d3b353ba6d4cfdac Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Dec 2021 08:12:37 -0500 Subject: [PATCH 01/17] Agent: Add IPayload interface --- monkey/infection_monkey/payload/i_payload.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 monkey/infection_monkey/payload/i_payload.py 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. + """ From c18af3c3fb37e1101f8e0a2d9c39b9c9fcd0271f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Dec 2021 08:14:53 -0500 Subject: [PATCH 02/17] Agent: Change return type of IPuppet.run_payload() to None At the moment, we don't expect payloads to return any values. This may be reevaluated as development proceeds or when telemetry is refactored. --- monkey/infection_monkey/i_puppet.py | 8 ++++---- monkey/infection_monkey/master/mock_master.py | 5 +---- monkey/infection_monkey/puppet/mock_puppet.py | 7 ++----- monkey/infection_monkey/puppet/puppet.py | 6 ++---- 4 files changed, 9 insertions(+), 17 deletions(-) diff --git a/monkey/infection_monkey/i_puppet.py b/monkey/infection_monkey/i_puppet.py index e25d20f53..50e050dc6 100644 --- a/monkey/infection_monkey/i_puppet.py +++ b/monkey/infection_monkey/i_puppet.py @@ -2,7 +2,7 @@ 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 @@ -107,13 +107,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/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/puppet/mock_puppet.py b/monkey/infection_monkey/puppet/mock_puppet.py index 204e44ab4..59539c58b 100644 --- a/monkey/infection_monkey/puppet/mock_puppet.py +++ b/monkey/infection_monkey/puppet/mock_puppet.py @@ -1,6 +1,6 @@ import logging import threading -from typing import Dict, Tuple +from typing import Dict from infection_monkey.i_puppet import ( ExploiterResultData, @@ -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/puppet.py b/monkey/infection_monkey/puppet/puppet.py index f932d84a4..5563a2dfe 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -1,6 +1,6 @@ import logging import threading -from typing import Dict, Tuple +from typing import Dict from infection_monkey.i_puppet import ( ExploiterResultData, @@ -45,9 +45,7 @@ class Puppet(IPuppet): ) -> ExploiterResultData: pass - 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): pass def cleanup(self) -> None: From 09a1297f47ccefa981406333f64053f6608c9f20 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Dec 2021 08:19:14 -0500 Subject: [PATCH 03/17] Agent: User relative imports within ransomware package --- .../infection_monkey/ransomware/file_selectors.py | 3 ++- .../ransomware/ransomware_payload.py | 5 +++-- .../ransomware/ransomware_payload_builder.py | 13 +++++++------ 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/monkey/infection_monkey/ransomware/file_selectors.py b/monkey/infection_monkey/ransomware/file_selectors.py index 33b73dd06..5707fba7d 100644 --- a/monkey/infection_monkey/ransomware/file_selectors.py +++ b/monkey/infection_monkey/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/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index a1e052970..227829d10 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -2,11 +2,12 @@ import logging 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_config import RansomwareConfig + logger = logging.getLogger(__name__) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload_builder.py b/monkey/infection_monkey/ransomware/ransomware_payload_builder.py index 9f0d78754..c5eb034f0 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload_builder.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload_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,6 +9,13 @@ 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_config import RansomwareConfig +from .ransomware_payload import RansomwarePayload +from .targeted_file_extensions import TARGETED_FILE_EXTENSIONS + EXTENSION = ".m0nk3y" CHUNK_SIZE = 4096 * 24 From 33e3a31030a0340a299ef9a6099cde9c052cc9bb Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Dec 2021 08:22:17 -0500 Subject: [PATCH 04/17] Agent: Move ransomware/ to payload/ransomware/ --- .gitattributes | 2 +- .../infection_monkey/{ => payload}/ransomware/__init__.py | 0 monkey/infection_monkey/{ => payload}/ransomware/consts.py | 0 .../{ => payload}/ransomware/file_selectors.py | 0 .../{ => payload}/ransomware/in_place_file_encryptor.py | 0 .../{ => payload}/ransomware/ransomware_config.py | 0 .../{ => payload}/ransomware/ransomware_payload.py | 0 .../{ => payload}/ransomware/ransomware_payload_builder.py | 0 .../{ => payload}/ransomware/ransomware_readme.txt | 0 .../{ => payload}/ransomware/readme_dropper.py | 0 .../{ => payload}/ransomware/targeted_file_extensions.py | 0 .../infection_monkey/ransomware/test_file_selectors.py | 4 ++-- .../ransomware/test_in_place_file_encryptor.py | 2 +- .../infection_monkey/ransomware/test_ransomware_config.py | 4 ++-- .../infection_monkey/ransomware/test_ransomware_payload.py | 6 +++--- .../infection_monkey/ransomware/test_readme_dropper.py | 2 +- 16 files changed, 10 insertions(+), 10 deletions(-) rename monkey/infection_monkey/{ => payload}/ransomware/__init__.py (100%) rename monkey/infection_monkey/{ => payload}/ransomware/consts.py (100%) rename monkey/infection_monkey/{ => payload}/ransomware/file_selectors.py (100%) rename monkey/infection_monkey/{ => payload}/ransomware/in_place_file_encryptor.py (100%) rename monkey/infection_monkey/{ => payload}/ransomware/ransomware_config.py (100%) rename monkey/infection_monkey/{ => payload}/ransomware/ransomware_payload.py (100%) rename monkey/infection_monkey/{ => payload}/ransomware/ransomware_payload_builder.py (100%) rename monkey/infection_monkey/{ => payload}/ransomware/ransomware_readme.txt (100%) rename monkey/infection_monkey/{ => payload}/ransomware/readme_dropper.py (100%) rename monkey/infection_monkey/{ => payload}/ransomware/targeted_file_extensions.py (100%) 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/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 100% rename from monkey/infection_monkey/ransomware/file_selectors.py rename to monkey/infection_monkey/payload/ransomware/file_selectors.py 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_config.py b/monkey/infection_monkey/payload/ransomware/ransomware_config.py similarity index 100% rename from monkey/infection_monkey/ransomware/ransomware_config.py rename to monkey/infection_monkey/payload/ransomware/ransomware_config.py diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/payload/ransomware/ransomware_payload.py similarity index 100% rename from monkey/infection_monkey/ransomware/ransomware_payload.py rename to monkey/infection_monkey/payload/ransomware/ransomware_payload.py diff --git a/monkey/infection_monkey/ransomware/ransomware_payload_builder.py b/monkey/infection_monkey/payload/ransomware/ransomware_payload_builder.py similarity index 100% rename from monkey/infection_monkey/ransomware/ransomware_payload_builder.py rename to monkey/infection_monkey/payload/ransomware/ransomware_payload_builder.py 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/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py index 42e852b95..23c723f0e 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py +++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_file_selectors.py @@ -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_payload 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/ransomware/test_in_place_file_encryptor.py index eb2633226..318536de9 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_in_place_file_encryptor.py +++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_in_place_file_encryptor.py @@ -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/ransomware/test_ransomware_config.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py index 141186f18..3d016f80c 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py +++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py @@ -4,8 +4,8 @@ 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 +from infection_monkey.payload.ransomware import ransomware_config +from infection_monkey.payload.ransomware.ransomware_config import RansomwareConfig LINUX_DIR = "/tmp/test" WINDOWS_DIR = "C:\\tmp\\test" diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py index 24eb8443d..a2b5babdb 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py +++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py @@ -7,9 +7,9 @@ from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import 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 +from infection_monkey.payload.ransomware.consts import README_FILE_NAME, README_SRC +from infection_monkey.payload.ransomware.ransomware_config import RansomwareConfig +from infection_monkey.payload.ransomware.ransomware_payload import RansomwarePayload @pytest.fixture diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_readme_dropper.py b/monkey/tests/unit_tests/infection_monkey/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/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" From ee1fa01dda9f09fbcfd615f4810b7fa76a32662e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Dec 2021 08:27:25 -0500 Subject: [PATCH 05/17] UT: Move ransomware unit tests to payload/ransomware/ --- .../infection_monkey/{ => payload}/ransomware/conftest.py | 0 .../{ => payload}/ransomware/ransomware_target_files.py | 0 .../{ => payload}/ransomware/test_file_selectors.py | 2 +- .../{ => payload}/ransomware/test_in_place_file_encryptor.py | 2 +- .../{ => payload}/ransomware/test_ransomware_config.py | 0 .../{ => payload}/ransomware/test_ransomware_payload.py | 2 +- .../{ => payload}/ransomware/test_readme_dropper.py | 0 7 files changed, 3 insertions(+), 3 deletions(-) rename monkey/tests/unit_tests/infection_monkey/{ => payload}/ransomware/conftest.py (100%) rename monkey/tests/unit_tests/infection_monkey/{ => payload}/ransomware/ransomware_target_files.py (100%) rename monkey/tests/unit_tests/infection_monkey/{ => payload}/ransomware/test_file_selectors.py (96%) rename monkey/tests/unit_tests/infection_monkey/{ => payload}/ransomware/test_in_place_file_encryptor.py (96%) rename monkey/tests/unit_tests/infection_monkey/{ => payload}/ransomware/test_ransomware_config.py (100%) rename monkey/tests/unit_tests/infection_monkey/{ => payload}/ransomware/test_ransomware_payload.py (98%) rename monkey/tests/unit_tests/infection_monkey/{ => payload}/ransomware/test_readme_dropper.py (100%) 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 96% 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 23c723f0e..635b39587 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, 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 96% 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 318536de9..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, diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_config.py similarity index 100% rename from monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py rename to monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_config.py diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_payload.py similarity index 98% rename from monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py rename to monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_payload.py index a2b5babdb..39ca057ee 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py +++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_payload.py @@ -2,7 +2,7 @@ from pathlib import PurePosixPath from unittest.mock import MagicMock 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, TEST_KEYBOARD_TXT, ) 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 100% 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 From 8e6abcb7956ffd1deb68616be8a93f137acfebfb Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 14 Dec 2021 20:40:34 +0530 Subject: [PATCH 06/17] Agent: Add PluginRegistry --- .../puppet/plugin_registry.py | 42 +++++++++++++++++++ monkey/infection_monkey/puppet/puppet.py | 6 ++- 2 files changed, 47 insertions(+), 1 deletion(-) create mode 100644 monkey/infection_monkey/puppet/plugin_registry.py diff --git a/monkey/infection_monkey/puppet/plugin_registry.py b/monkey/infection_monkey/puppet/plugin_registry.py new file mode 100644 index 000000000..d86a7491e --- /dev/null +++ b/monkey/infection_monkey/puppet/plugin_registry.py @@ -0,0 +1,42 @@ +import logging +from typing import Optional + +from infection_monkey.i_puppet import UnknownPluginError +from infection_monkey.puppet.plugin_type import PluginType + +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: object, plugin_type: PluginType) -> None: + self._registry.setdefault(plugin_type, {}) + self._registry[plugin_type][plugin.__class__.__name__] = plugin + + logger.debug(f"Plugin '{plugin.__class__.__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 5563a2dfe..0c36435e0 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -10,14 +10,18 @@ from infection_monkey.i_puppet import ( PortScanData, PostBreachData, ) +from infection_monkey.puppet.plugin_registry import PluginRegistry from infection_monkey.puppet.plugin_type import PluginType logger = logging.getLogger() class Puppet(IPuppet): + def __init__(self) -> None: + self._plugin_registry = PluginRegistry() + def load_plugin(self, plugin: object, plugin_type: PluginType) -> None: - pass + self._plugin_registry.load_plugin(plugin, plugin_type) def run_sys_info_collector(self, name: str) -> Dict: pass From b7982552498b3f7fbcc25edf3ae94471f0fe41c1 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 16 Dec 2021 16:26:40 +0100 Subject: [PATCH 07/17] Agent: Add plugin_name attribute to puppet's load_plugin --- monkey/infection_monkey/i_puppet.py | 5 +++-- monkey/infection_monkey/puppet/plugin_registry.py | 6 +++--- monkey/infection_monkey/puppet/puppet.py | 4 ++-- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/i_puppet.py b/monkey/infection_monkey/i_puppet.py index 50e050dc6..f3567ee77 100644 --- a/monkey/infection_monkey/i_puppet.py +++ b/monkey/infection_monkey/i_puppet.py @@ -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 """ diff --git a/monkey/infection_monkey/puppet/plugin_registry.py b/monkey/infection_monkey/puppet/plugin_registry.py index d86a7491e..76d0d3714 100644 --- a/monkey/infection_monkey/puppet/plugin_registry.py +++ b/monkey/infection_monkey/puppet/plugin_registry.py @@ -23,11 +23,11 @@ class PluginRegistry: """ self._registry = {} - def load_plugin(self, plugin: object, plugin_type: PluginType) -> None: + def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None: self._registry.setdefault(plugin_type, {}) - self._registry[plugin_type][plugin.__class__.__name__] = plugin + self._registry[plugin_type][plugin_name] = plugin - logger.debug(f"Plugin '{plugin.__class__.__name__}' loaded") + logger.debug(f"Plugin '{plugin_name}' loaded") def get_plugin(self, plugin_name: str, plugin_type: PluginType) -> Optional[object]: try: diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 0c36435e0..0fac17a7e 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -20,8 +20,8 @@ class Puppet(IPuppet): def __init__(self) -> None: self._plugin_registry = PluginRegistry() - def load_plugin(self, plugin: object, plugin_type: PluginType) -> None: - self._plugin_registry.load_plugin(plugin, plugin_type) + def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None: + self._plugin_registry.load_plugin(plugin, plugin_name, plugin_type) def run_sys_info_collector(self, name: str) -> Dict: pass From 0a4ff25843d5d9fab249db94e1f583bd96c6107b Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 16 Dec 2021 16:28:46 +0100 Subject: [PATCH 08/17] Agent: Implement Puppet.run_payload() --- monkey/infection_monkey/puppet/puppet.py | 5 ++- .../infection_monkey/puppet/test_puppet.py | 43 +++++++++++++++++++ 2 files changed, 46 insertions(+), 2 deletions(-) create mode 100644 monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 0fac17a7e..132898536 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -21,7 +21,7 @@ class Puppet(IPuppet): self._plugin_registry = PluginRegistry() def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None: - self._plugin_registry.load_plugin(plugin, plugin_name, plugin_type) + self._plugin_registry.load_plugin(plugin_name, plugin, plugin_type) def run_sys_info_collector(self, name: str) -> Dict: pass @@ -50,7 +50,8 @@ class Puppet(IPuppet): pass def run_payload(self, name: str, options: Dict, interrupt: threading.Event): - pass + 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/puppet/test_puppet.py b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py new file mode 100644 index 000000000..c0c4a1f19 --- /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.puppet.plugin_type 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() From 2299c029d7c36706990bd2716db6bbc7866fe6ec Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Dec 2021 09:08:30 -0500 Subject: [PATCH 09/17] Agent: Rename RansomwarePayload to Ransomware A payload adheres to a specific IPayload interface. The class that is now called RansomwarePayload is just a concrete ransomware. A new RansomwarePayload will be introduced to wrap the build and execute of the Ransomware. --- .../{ransomware_payload.py => ransomware.py} | 2 +- ...yload_builder.py => ransomware_builder.py} | 8 +- .../payload/ransomware/test_file_selectors.py | 2 +- ...nsomware_payload.py => test_ransomware.py} | 98 ++++++++----------- 4 files changed, 49 insertions(+), 61 deletions(-) rename monkey/infection_monkey/payload/ransomware/{ransomware_payload.py => ransomware.py} (99%) rename monkey/infection_monkey/payload/ransomware/{ransomware_payload_builder.py => ransomware_builder.py} (89%) rename monkey/tests/unit_tests/infection_monkey/payload/ransomware/{test_ransomware_payload.py => test_ransomware.py} (55%) diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_payload.py b/monkey/infection_monkey/payload/ransomware/ransomware.py similarity index 99% rename from monkey/infection_monkey/payload/ransomware/ransomware_payload.py rename to monkey/infection_monkey/payload/ransomware/ransomware.py index 227829d10..2f09e386f 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware.py @@ -11,7 +11,7 @@ from .ransomware_config import RansomwareConfig logger = logging.getLogger(__name__) -class RansomwarePayload: +class Ransomware: def __init__( self, config: RansomwareConfig, diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_payload_builder.py b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py similarity index 89% rename from monkey/infection_monkey/payload/ransomware/ransomware_payload_builder.py rename to monkey/infection_monkey/payload/ransomware/ransomware_builder.py index c5eb034f0..671f20f1f 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware_payload_builder.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py @@ -12,8 +12,8 @@ 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_config import RansomwareConfig -from .ransomware_payload import RansomwarePayload from .targeted_file_extensions import TARGETED_FILE_EXTENSIONS EXTENSION = ".m0nk3y" @@ -22,8 +22,8 @@ CHUNK_SIZE = 4096 * 24 logger = logging.getLogger(__name__) -def build_ransomware_payload(config: dict): - logger.debug(f"Ransomware payload configuration:\n{pformat(config)}") +def build_ransomware(config: dict): + logger.debug(f"Ransomware configuration:\n{pformat(config)}") ransomware_config = RansomwareConfig(config) file_encryptor = _build_file_encryptor() @@ -31,7 +31,7 @@ def build_ransomware_payload(config: dict): leave_readme = _build_leave_readme() telemetry_messenger = _build_telemetry_messenger() - return RansomwarePayload( + return Ransomware( ransomware_config, file_encryptor, file_selector, diff --git a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_file_selectors.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_file_selectors.py index 635b39587..f779b733e 100644 --- a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_file_selectors.py +++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_file_selectors.py @@ -13,7 +13,7 @@ from tests.unit_tests.infection_monkey.payload.ransomware.ransomware_target_file from tests.utils import is_user_admin from infection_monkey.payload.ransomware.file_selectors import ProductionSafeTargetFileSelector -from infection_monkey.payload.ransomware.ransomware_payload import README_SRC +from infection_monkey.payload.ransomware.ransomware import README_SRC TARGETED_FILE_EXTENSIONS = [".pdf", ".txt"] diff --git a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_payload.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py similarity index 55% rename from monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_payload.py rename to monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py index 39ca057ee..a7e9f8a90 100644 --- a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_payload.py +++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py @@ -8,17 +8,17 @@ from tests.unit_tests.infection_monkey.payload.ransomware.ransomware_target_file ) 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_config import RansomwareConfig -from infection_monkey.payload.ransomware.ransomware_payload import RansomwarePayload @pytest.fixture -def ransomware_payload(build_ransomware_payload, ransomware_payload_config): - return build_ransomware_payload(ransomware_payload_config) +def ransomware(build_ransomware, ransomware_config): + return build_ransomware(ransomware_config) @pytest.fixture -def build_ransomware_payload( +def build_ransomware( mock_file_encryptor, mock_file_selector, mock_leave_readme, telemetry_messenger_spy ): def inner( @@ -27,7 +27,7 @@ def build_ransomware_payload( file_selector=mock_file_selector, leave_readme=mock_leave_readme, ): - return RansomwarePayload( + return Ransomware( config, file_encryptor, file_selector, @@ -39,7 +39,7 @@ def build_ransomware_payload( @pytest.fixture -def ransomware_payload_config(ransomware_test_data): +def ransomware_config(ransomware_test_data): class RansomwareConfigStub(RansomwareConfig): def __init__(self, encryption_enabled, readme_enabled, target_directory): self.encryption_enabled = encryption_enabled @@ -69,18 +69,16 @@ def mock_leave_readme(): def test_files_selected_from_target_dir( - ransomware_payload, - ransomware_payload_config, + ransomware, + ransomware_config, mock_file_selector, ): - ransomware_payload.run_payload() - mock_file_selector.assert_called_with(ransomware_payload_config.target_directory) + ransomware.run_payload() + mock_file_selector.assert_called_with(ransomware_config.target_directory) -def test_all_selected_files_encrypted( - ransomware_test_data, ransomware_payload, mock_file_encryptor -): - ransomware_payload.run_payload() +def test_all_selected_files_encrypted(ransomware_test_data, ransomware, mock_file_encryptor): + ransomware.run_payload() assert mock_file_encryptor.call_count == 2 mock_file_encryptor.assert_any_call(ransomware_test_data / ALL_ZEROS_PDF) @@ -88,30 +86,30 @@ def test_all_selected_files_encrypted( def test_encryption_skipped_if_configured_false( - build_ransomware_payload, ransomware_payload_config, mock_file_encryptor + build_ransomware, ransomware_config, mock_file_encryptor ): - ransomware_payload_config.encryption_enabled = False + ransomware_config.encryption_enabled = False - ransomware_payload = build_ransomware_payload(ransomware_payload_config) - ransomware_payload.run_payload() + ransomware = build_ransomware(ransomware_config) + ransomware.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 + build_ransomware, ransomware_config, mock_file_encryptor ): - ransomware_payload_config.encryption_enabled = True - ransomware_payload_config.target_directory = None + ransomware_config.encryption_enabled = True + ransomware_config.target_directory = None - ransomware_payload = build_ransomware_payload(ransomware_payload_config) - ransomware_payload.run_payload() + ransomware = build_ransomware(ransomware_config) + ransomware.run_payload() assert mock_file_encryptor.call_count == 0 -def test_telemetry_success(ransomware_payload, telemetry_messenger_spy): - ransomware_payload.run_payload() +def test_telemetry_success(ransomware, telemetry_messenger_spy): + ransomware.run_payload() assert len(telemetry_messenger_spy.telemetries) == 2 telem_1 = telemetry_messenger_spy.telemetries[0] @@ -125,19 +123,15 @@ def test_telemetry_success(ransomware_payload, telemetry_messenger_spy): assert telem_2.get_data()["files"][0]["error"] == "" -def test_telemetry_failure( - build_ransomware_payload, ransomware_payload_config, telemetry_messenger_spy -): +def test_telemetry_failure(build_ransomware, ransomware_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 = build_ransomware(config=ransomware_config, file_encryptor=mfe, file_selector=mfs) - ransomware_payload.run_payload() + ransomware.run_payload() telem = telemetry_messenger_spy.telemetries[0] assert file_not_exists in telem.get_data()["files"][0]["path"] @@ -145,42 +139,36 @@ def test_telemetry_failure( 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) +def test_readme_false(build_ransomware, ransomware_config, mock_leave_readme): + ransomware_config.readme_enabled = False + ransomware = build_ransomware(ransomware_config) - ransomware_payload.run_payload() + ransomware.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) +def test_readme_true(build_ransomware, ransomware_config, mock_leave_readme, ransomware_test_data): + ransomware_config.readme_enabled = True + ransomware = build_ransomware(ransomware_config) - ransomware_payload.run_payload() + ransomware.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 +def test_no_readme_if_no_directory(build_ransomware, ransomware_config, mock_leave_readme): + ransomware_config.target_directory = None + ransomware_config.readme_enabled = True - ransomware_payload = build_ransomware_payload(ransomware_payload_config) + ransomware = build_ransomware(ransomware_config) - ransomware_payload.run_payload() + ransomware.run_payload() mock_leave_readme.assert_not_called() -def test_leave_readme_exceptions_handled(build_ransomware_payload, ransomware_payload_config): +def test_leave_readme_exceptions_handled(build_ransomware, ransomware_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 - ) + ransomware_config.readme_enabled = True + ransomware = build_ransomware(config=ransomware_config, leave_readme=leave_readme) # Test will fail if exception is raised and not handled - ransomware_payload.run_payload() + ransomware.run_payload() From 0328d2860eb8af12ff40157853d68637bb7cffdd Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Dec 2021 09:17:19 -0500 Subject: [PATCH 10/17] Agent: Add a RansomwarePayload that implements to the IPayload interface --- .../payload/ransomware/ransomware.py | 3 ++- .../payload/ransomware/ransomware_payload.py | 12 +++++++++++ .../payload/ransomware/test_ransomware.py | 21 ++++++++++--------- 3 files changed, 25 insertions(+), 11 deletions(-) create mode 100644 monkey/infection_monkey/payload/ransomware/ransomware_payload.py diff --git a/monkey/infection_monkey/payload/ransomware/ransomware.py b/monkey/infection_monkey/payload/ransomware/ransomware.py index 2f09e386f..1050bab75 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware.py @@ -1,4 +1,5 @@ import logging +import threading from pathlib import Path from typing import Callable, List @@ -32,7 +33,7 @@ class Ransomware: self._target_directory / README_FILE_NAME if self._target_directory else None ) - def run_payload(self): + def run(self, _: threading.Event): if not self._target_directory: return 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/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py index a7e9f8a90..6024f2afd 100644 --- a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py +++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py @@ -1,3 +1,4 @@ +import threading from pathlib import PurePosixPath from unittest.mock import MagicMock @@ -73,12 +74,12 @@ def test_files_selected_from_target_dir( ransomware_config, mock_file_selector, ): - ransomware.run_payload() + ransomware.run(threading.Event()) mock_file_selector.assert_called_with(ransomware_config.target_directory) def test_all_selected_files_encrypted(ransomware_test_data, ransomware, mock_file_encryptor): - ransomware.run_payload() + ransomware.run(threading.Event()) assert mock_file_encryptor.call_count == 2 mock_file_encryptor.assert_any_call(ransomware_test_data / ALL_ZEROS_PDF) @@ -91,7 +92,7 @@ def test_encryption_skipped_if_configured_false( ransomware_config.encryption_enabled = False ransomware = build_ransomware(ransomware_config) - ransomware.run_payload() + ransomware.run(threading.Event()) assert mock_file_encryptor.call_count == 0 @@ -103,13 +104,13 @@ def test_encryption_skipped_if_no_directory( ransomware_config.target_directory = None ransomware = build_ransomware(ransomware_config) - ransomware.run_payload() + ransomware.run(threading.Event()) assert mock_file_encryptor.call_count == 0 def test_telemetry_success(ransomware, telemetry_messenger_spy): - ransomware.run_payload() + ransomware.run(threading.Event()) assert len(telemetry_messenger_spy.telemetries) == 2 telem_1 = telemetry_messenger_spy.telemetries[0] @@ -131,7 +132,7 @@ def test_telemetry_failure(build_ransomware, ransomware_config, telemetry_messen mfs = MagicMock(return_value=[PurePosixPath(file_not_exists)]) ransomware = build_ransomware(config=ransomware_config, file_encryptor=mfe, file_selector=mfs) - ransomware.run_payload() + ransomware.run(threading.Event()) telem = telemetry_messenger_spy.telemetries[0] assert file_not_exists in telem.get_data()["files"][0]["path"] @@ -143,7 +144,7 @@ def test_readme_false(build_ransomware, ransomware_config, mock_leave_readme): ransomware_config.readme_enabled = False ransomware = build_ransomware(ransomware_config) - ransomware.run_payload() + ransomware.run(threading.Event()) mock_leave_readme.assert_not_called() @@ -151,7 +152,7 @@ def test_readme_true(build_ransomware, ransomware_config, mock_leave_readme, ran ransomware_config.readme_enabled = True ransomware = build_ransomware(ransomware_config) - ransomware.run_payload() + ransomware.run(threading.Event()) mock_leave_readme.assert_called_with(README_SRC, ransomware_test_data / README_FILE_NAME) @@ -161,7 +162,7 @@ def test_no_readme_if_no_directory(build_ransomware, ransomware_config, mock_lea ransomware = build_ransomware(ransomware_config) - ransomware.run_payload() + ransomware.run(threading.Event()) mock_leave_readme.assert_not_called() @@ -171,4 +172,4 @@ def test_leave_readme_exceptions_handled(build_ransomware, ransomware_config): ransomware = build_ransomware(config=ransomware_config, leave_readme=leave_readme) # Test will fail if exception is raised and not handled - ransomware.run_payload() + ransomware.run(threading.Event()) From 958cf3a25298a77b173fb535ad55981c1008744d Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 17 Dec 2021 19:55:26 +0530 Subject: [PATCH 11/17] Agent, UT: Rename 'config' to 'options' in ransomware files --- .../payload/ransomware/ransomware.py | 4 +- .../payload/ransomware/ransomware_builder.py | 10 +-- ...omware_config.py => ransomware_options.py} | 10 +-- .../payload/ransomware/test_ransomware.py | 60 +++++++-------- .../ransomware/test_ransomware_config.py | 73 ------------------- .../ransomware/test_ransomware_options.py | 73 +++++++++++++++++++ .../monkey_island/cc/services/test_config.py | 4 +- 7 files changed, 117 insertions(+), 117 deletions(-) rename monkey/infection_monkey/payload/ransomware/{ransomware_config.py => ransomware_options.py} (72%) delete mode 100644 monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_config.py create mode 100644 monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_options.py diff --git a/monkey/infection_monkey/payload/ransomware/ransomware.py b/monkey/infection_monkey/payload/ransomware/ransomware.py index 1050bab75..d83361dca 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware.py @@ -7,7 +7,7 @@ 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_config import RansomwareConfig +from .ransomware_options import RansomwareOptions logger = logging.getLogger(__name__) @@ -15,7 +15,7 @@ logger = logging.getLogger(__name__) 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], diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_builder.py b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py index 671f20f1f..4b8bbc8bb 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware_builder.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py @@ -13,7 +13,7 @@ from . import readme_dropper from .file_selectors import ProductionSafeTargetFileSelector from .in_place_file_encryptor import InPlaceFileEncryptor from .ransomware import Ransomware -from .ransomware_config import RansomwareConfig +from .ransomware_options import RansomwareOptions from .targeted_file_extensions import TARGETED_FILE_EXTENSIONS EXTENSION = ".m0nk3y" @@ -22,9 +22,9 @@ CHUNK_SIZE = 4096 * 24 logger = logging.getLogger(__name__) -def build_ransomware(config: dict): - logger.debug(f"Ransomware 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() @@ -32,7 +32,7 @@ def build_ransomware(config: dict): telemetry_messenger = _build_telemetry_messenger() return Ransomware( - ransomware_config, + ransomware_options, file_encryptor, file_selector, leave_readme, diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_config.py b/monkey/infection_monkey/payload/ransomware/ransomware_options.py similarity index 72% rename from monkey/infection_monkey/payload/ransomware/ransomware_config.py rename to monkey/infection_monkey/payload/ransomware/ransomware_options.py index f8ab792da..8416f8465 100644 --- a/monkey/infection_monkey/payload/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/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py index 6024f2afd..94de5aabc 100644 --- a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py +++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py @@ -10,12 +10,12 @@ from tests.unit_tests.infection_monkey.payload.ransomware.ransomware_target_file 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_config import RansomwareConfig +from infection_monkey.payload.ransomware.ransomware_options import RansomwareOptions @pytest.fixture -def ransomware(build_ransomware, ransomware_config): - return build_ransomware(ransomware_config) +def ransomware(build_ransomware, ransomware_options): + return build_ransomware(ransomware_options) @pytest.fixture @@ -40,14 +40,14 @@ def build_ransomware( @pytest.fixture -def ransomware_config(ransomware_test_data): - class RansomwareConfigStub(RansomwareConfig): +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 RansomwareConfigStub(True, False, ransomware_test_data) + return RansomwareOptionsStub(True, False, ransomware_test_data) @pytest.fixture @@ -71,11 +71,11 @@ def mock_leave_readme(): def test_files_selected_from_target_dir( ransomware, - ransomware_config, + ransomware_options, mock_file_selector, ): ransomware.run(threading.Event()) - mock_file_selector.assert_called_with(ransomware_config.target_directory) + mock_file_selector.assert_called_with(ransomware_options.target_directory) def test_all_selected_files_encrypted(ransomware_test_data, ransomware, mock_file_encryptor): @@ -87,23 +87,23 @@ def test_all_selected_files_encrypted(ransomware_test_data, ransomware, mock_fil def test_encryption_skipped_if_configured_false( - build_ransomware, ransomware_config, mock_file_encryptor + build_ransomware, ransomware_options, mock_file_encryptor ): - ransomware_config.encryption_enabled = False + ransomware_options.encryption_enabled = False - ransomware = build_ransomware(ransomware_config) + 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_config, mock_file_encryptor + build_ransomware, ransomware_options, mock_file_encryptor ): - ransomware_config.encryption_enabled = True - ransomware_config.target_directory = None + ransomware_options.encryption_enabled = True + ransomware_options.target_directory = None - ransomware = build_ransomware(ransomware_config) + ransomware = build_ransomware(ransomware_options) ransomware.run(threading.Event()) assert mock_file_encryptor.call_count == 0 @@ -124,13 +124,13 @@ def test_telemetry_success(ransomware, telemetry_messenger_spy): assert telem_2.get_data()["files"][0]["error"] == "" -def test_telemetry_failure(build_ransomware, ransomware_config, telemetry_messenger_spy): +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_config, file_encryptor=mfe, file_selector=mfs) + ransomware = build_ransomware(config=ransomware_options, file_encryptor=mfe, file_selector=mfs) ransomware.run(threading.Event()) telem = telemetry_messenger_spy.telemetries[0] @@ -140,36 +140,36 @@ def test_telemetry_failure(build_ransomware, ransomware_config, telemetry_messen assert "No such file or directory" in telem.get_data()["files"][0]["error"] -def test_readme_false(build_ransomware, ransomware_config, mock_leave_readme): - ransomware_config.readme_enabled = False - ransomware = build_ransomware(ransomware_config) +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_config, mock_leave_readme, ransomware_test_data): - ransomware_config.readme_enabled = True - ransomware = build_ransomware(ransomware_config) +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_config, mock_leave_readme): - ransomware_config.target_directory = None - ransomware_config.readme_enabled = True +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_config) + ransomware = build_ransomware(ransomware_options) ransomware.run(threading.Event()) mock_leave_readme.assert_not_called() -def test_leave_readme_exceptions_handled(build_ransomware, ransomware_config): +def test_leave_readme_exceptions_handled(build_ransomware, ransomware_options): leave_readme = MagicMock(side_effect=Exception("Test exception when leaving README")) - ransomware_config.readme_enabled = True - ransomware = build_ransomware(config=ransomware_config, leave_readme=leave_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_config.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_config.py deleted file mode 100644 index 3d016f80c..000000000 --- a/monkey/tests/unit_tests/infection_monkey/payload/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.payload.ransomware import ransomware_config -from infection_monkey.payload.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/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/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 From 61a7647f9bf0333acf845731b16659d313c31a37 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 17 Dec 2021 15:31:20 +0100 Subject: [PATCH 12/17] Agent: Add interrupt handling to ransomware --- .../payload/ransomware/ransomware.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/payload/ransomware/ransomware.py b/monkey/infection_monkey/payload/ransomware/ransomware.py index d83361dca..fcc0a53e7 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware.py @@ -33,7 +33,7 @@ class Ransomware: self._target_directory / README_FILE_NAME if self._target_directory else None ) - def run(self, _: threading.Event): + def run(self, interrupt: threading.Event): if not self._target_directory: return @@ -41,7 +41,11 @@ class Ransomware: if self._config.encryption_enabled: file_list = self._find_files() - self._encrypt_files(file_list) + self._encrypt_files(file_list, interrupt) + + if interrupt.is_set(): + logger.debug("Received a stop signal, skipping remaining tasks of ransomware payload") + return if self._config.readme_enabled: self._leave_readme_in_target_directory() @@ -50,10 +54,16 @@ class Ransomware: 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) From 05c57644874572acf48332a07ca2042c0f7ad79b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Dec 2021 09:40:16 -0500 Subject: [PATCH 13/17] Agent: Add i_puppet package --- monkey/infection_monkey/i_puppet/__init__.py | 10 ++++++++++ monkey/infection_monkey/{ => i_puppet}/i_puppet.py | 0 2 files changed, 10 insertions(+) create mode 100644 monkey/infection_monkey/i_puppet/__init__.py rename monkey/infection_monkey/{ => i_puppet}/i_puppet.py (100%) diff --git a/monkey/infection_monkey/i_puppet/__init__.py b/monkey/infection_monkey/i_puppet/__init__.py new file mode 100644 index 000000000..7140b9082 --- /dev/null +++ b/monkey/infection_monkey/i_puppet/__init__.py @@ -0,0 +1,10 @@ +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 100% rename from monkey/infection_monkey/i_puppet.py rename to monkey/infection_monkey/i_puppet/i_puppet.py From afbc313a7cff0d16e9b3c24cf33d2b9deba5a672 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 17 Dec 2021 16:10:42 +0100 Subject: [PATCH 14/17] Agent: Handle interrupts in ransomware --- .../payload/ransomware/ransomware.py | 12 +++--- .../payload/ransomware/test_ransomware.py | 38 +++++++++++++++++++ 2 files changed, 44 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/payload/ransomware/ransomware.py b/monkey/infection_monkey/payload/ransomware/ransomware.py index fcc0a53e7..003112cc3 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware.py @@ -43,12 +43,8 @@ class Ransomware: file_list = self._find_files() self._encrypt_files(file_list, interrupt) - if interrupt.is_set(): - logger.debug("Received a stop signal, skipping remaining tasks of ransomware payload") - return - 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}") @@ -76,8 +72,12 @@ class Ransomware: 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/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py index 94de5aabc..adffe6f88 100644 --- a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py +++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py @@ -5,6 +5,7 @@ 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, ) @@ -69,6 +70,11 @@ def mock_leave_readme(): return MagicMock() +@pytest.fixture +def interrupt(): + return threading.Event() + + def test_files_selected_from_target_dir( ransomware, ransomware_options, @@ -86,6 +92,38 @@ def test_all_selected_files_encrypted(ransomware_test_data, ransomware, mock_fil 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 ): From 973c88678ea08d3719bcdd0e8f749b1ee793c99d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Dec 2021 10:13:28 -0500 Subject: [PATCH 15/17] Agent: Move PluginType to the i_plugin package --- monkey/infection_monkey/i_puppet/__init__.py | 1 + monkey/infection_monkey/i_puppet/i_puppet.py | 2 +- monkey/infection_monkey/{puppet => i_puppet}/plugin_type.py | 0 monkey/infection_monkey/puppet/mock_puppet.py | 2 +- monkey/infection_monkey/puppet/plugin_registry.py | 3 +-- monkey/infection_monkey/puppet/puppet.py | 2 +- monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py | 2 +- 7 files changed, 6 insertions(+), 6 deletions(-) rename monkey/infection_monkey/{puppet => i_puppet}/plugin_type.py (100%) diff --git a/monkey/infection_monkey/i_puppet/__init__.py b/monkey/infection_monkey/i_puppet/__init__.py index 7140b9082..0ba1096d1 100644 --- a/monkey/infection_monkey/i_puppet/__init__.py +++ b/monkey/infection_monkey/i_puppet/__init__.py @@ -1,3 +1,4 @@ +from .plugin_type import PluginType from .i_puppet import ( IPuppet, ExploiterResultData, diff --git a/monkey/infection_monkey/i_puppet/i_puppet.py b/monkey/infection_monkey/i_puppet/i_puppet.py index f3567ee77..3fa2aabd9 100644 --- a/monkey/infection_monkey/i_puppet/i_puppet.py +++ b/monkey/infection_monkey/i_puppet/i_puppet.py @@ -4,7 +4,7 @@ from collections import namedtuple from enum import Enum from typing import Dict -from infection_monkey.puppet.plugin_type import PluginType +from . import PluginType class PortStatus(Enum): 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/puppet/mock_puppet.py b/monkey/infection_monkey/puppet/mock_puppet.py index 59539c58b..182ebe55e 100644 --- a/monkey/infection_monkey/puppet/mock_puppet.py +++ b/monkey/infection_monkey/puppet/mock_puppet.py @@ -7,11 +7,11 @@ from infection_monkey.i_puppet import ( 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" diff --git a/monkey/infection_monkey/puppet/plugin_registry.py b/monkey/infection_monkey/puppet/plugin_registry.py index 76d0d3714..0e98ba2ef 100644 --- a/monkey/infection_monkey/puppet/plugin_registry.py +++ b/monkey/infection_monkey/puppet/plugin_registry.py @@ -1,8 +1,7 @@ import logging from typing import Optional -from infection_monkey.i_puppet import UnknownPluginError -from infection_monkey.puppet.plugin_type import PluginType +from infection_monkey.i_puppet import PluginType, UnknownPluginError logger = logging.getLogger() diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 132898536..64a8028e3 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -7,11 +7,11 @@ from infection_monkey.i_puppet import ( FingerprintData, IPuppet, PingScanData, + PluginType, PortScanData, PostBreachData, ) from infection_monkey.puppet.plugin_registry import PluginRegistry -from infection_monkey.puppet.plugin_type import PluginType logger = logging.getLogger() diff --git a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py index c0c4a1f19..950bc329b 100644 --- a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py +++ b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py @@ -1,7 +1,7 @@ import threading from unittest.mock import MagicMock -from infection_monkey.puppet.plugin_type import PluginType +from infection_monkey.i_puppet import PluginType from infection_monkey.puppet.puppet import Puppet From 7b8b485b5749b7d5ff27e0f6aaeeb4b411afccf0 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Dec 2021 10:22:42 -0500 Subject: [PATCH 16/17] Agent: Mock out unimplemented functions in Puppet --- monkey/infection_monkey/puppet/puppet.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 64a8028e3..3e2ad0f5f 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -13,27 +13,30 @@ from infection_monkey.i_puppet import ( ) from infection_monkey.puppet.plugin_registry import PluginRegistry +from .mock_puppet import MockPuppet + logger = logging.getLogger() class Puppet(IPuppet): 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, @@ -42,12 +45,12 @@ 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): payload = self._plugin_registry.get_plugin(name, PluginType.PAYLOAD) From b19ce79df678cfbbecfd96574ff4f708027a1716 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Dec 2021 10:25:16 -0500 Subject: [PATCH 17/17] Agent: Use relative imports within puppet package --- monkey/infection_monkey/puppet/puppet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 3e2ad0f5f..41ed99250 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -11,9 +11,9 @@ from infection_monkey.i_puppet import ( PortScanData, PostBreachData, ) -from infection_monkey.puppet.plugin_registry import PluginRegistry from .mock_puppet import MockPuppet +from .plugin_registry import PluginRegistry logger = logging.getLogger()