From 24915ba797f27d48fa54b2ce46c510fdd5df40c3 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 30 Mar 2022 14:13:51 +0300 Subject: [PATCH] Agent: Load and fix the custom PBA into puppet --- monkey/infection_monkey/i_puppet/i_puppet.py | 10 +-------- .../master/automated_master.py | 13 +++-------- monkey/infection_monkey/monkey.py | 4 ++++ .../custom_pba/users_custom_pba.py | 22 +++++++++++++------ monkey/monkey_island/cc/services/config.py | 2 ++ .../infection_monkey/master/mock_puppet.py | 5 +---- 6 files changed, 26 insertions(+), 30 deletions(-) diff --git a/monkey/infection_monkey/i_puppet/i_puppet.py b/monkey/infection_monkey/i_puppet/i_puppet.py index fd42d5c58..c4f46d792 100644 --- a/monkey/infection_monkey/i_puppet/i_puppet.py +++ b/monkey/infection_monkey/i_puppet/i_puppet.py @@ -3,7 +3,7 @@ import threading from collections import namedtuple from dataclasses import dataclass from enum import Enum -from typing import Any, Dict, Iterable, List, Mapping, Sequence +from typing import Dict, Iterable, List, Mapping, Sequence from infection_monkey.model import VictimHost @@ -67,14 +67,6 @@ class IPuppet(metaclass=abc.ABCMeta): :rtype: Iterable[PostBreachData] """ - @abc.abstractmethod - def run_custom_pba(self, options: Mapping[str, Any]) -> PostBreachData: - """ - Runs a user configured post breach action (PBA) - :param Dict options: A dictionary containing options that modify the behavior of the PBA - :rtype: PostBreachData - """ - @abc.abstractmethod def ping(self, host: str, timeout: float) -> PingScanData: """ diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 42086fee4..c76a8938c 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -4,7 +4,6 @@ import time from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple from infection_monkey.credential_store import ICredentialsStore -from common.common_consts.post_breach_consts import POST_BREACH_FILE_EXECUTION from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError from infection_monkey.i_master import IMaster from infection_monkey.i_puppet import IPuppet @@ -155,7 +154,7 @@ class AutomatedMaster(IMaster): ), ) pba_thread = create_daemon_thread( - target=self._run_PBAs, + target=self._run_pbas, name="PBAThread", args=(config["post_breach_actions"].items(), self._run_pba, config["custom_pbas"]), ) @@ -197,10 +196,6 @@ class AutomatedMaster(IMaster): name = pba[0] options = pba[1] - # TEMPORARY; TO AVOID ERRORS SINCE THIS ISN'T IMPLEMENTED YET - if name == "Custom": - return - for pba_data in self._puppet.run_pba(name, options): self._telemetry_messenger.send_telemetry(PostBreachTelem(pba_data)) @@ -213,14 +208,12 @@ class AutomatedMaster(IMaster): self._puppet.run_payload(name, options, self._stop) - def _run_PBAs( + def _run_pbas( self, plugins: Iterable[Any], callback: Callable[[Any], None], custom_pba_options: Mapping ): self._run_plugins(plugins, "post-breach action", callback) - command, result = self._puppet.run_custom_pba(custom_pba_options) - telem = PostBreachTelem(POST_BREACH_FILE_EXECUTION, command, result) - self._telemetry_messenger.send_telemetry(telem) + self._run_plugins([("CustomPBA", custom_pba_options)], "post-breach action", callback) def _run_plugins( self, plugins: Iterable[Any], plugin_type: str, callback: Callable[[Any], None] diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index f9eac872d..1863da03f 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -51,6 +51,7 @@ from infection_monkey.post_breach.actions.schedule_jobs import ScheduleJobs from infection_monkey.post_breach.actions.timestomping import Timestomping from infection_monkey.post_breach.actions.use_signed_scripts import SignedScriptProxyExecution from infection_monkey.post_breach.actions.use_trap_command import TrapCommand +from infection_monkey.post_breach.custom_pba.users_custom_pba import UsersPBA from infection_monkey.puppet.puppet import Puppet from infection_monkey.system_singleton import SystemSingleton from infection_monkey.telemetry.attack.t1106_telem import T1106Telem @@ -315,6 +316,9 @@ class InfectionMonkey: ClearCommandHistory(self._telemetry_messenger), PluginType.POST_BREACH_ACTION, ) + puppet.load_plugin( + "CustomPBA", UsersPBA(self._telemetry_messenger), PluginType.POST_BREACH_ACTION + ) puppet.load_plugin("ransomware", RansomwarePayload(), PluginType.PAYLOAD) diff --git a/monkey/infection_monkey/post_breach/custom_pba/users_custom_pba.py b/monkey/infection_monkey/post_breach/custom_pba/users_custom_pba.py index 81b6e9ab2..5f9100d98 100644 --- a/monkey/infection_monkey/post_breach/custom_pba/users_custom_pba.py +++ b/monkey/infection_monkey/post_breach/custom_pba/users_custom_pba.py @@ -1,11 +1,11 @@ import logging import os -from typing import Any, Mapping +from typing import Dict, Iterable from common.common_consts.post_breach_consts import POST_BREACH_FILE_EXECUTION from common.utils.attack_utils import ScanStatus -from infection_monkey.config import WormConfiguration from infection_monkey.control import ControlClient +from infection_monkey.i_puppet import PostBreachData from infection_monkey.network.tools import get_interface_to_target from infection_monkey.post_breach.pba import PBA from infection_monkey.telemetry.attack.t1105_telem import T1105Telem @@ -25,10 +25,18 @@ class UsersPBA(PBA): Defines user's configured post breach action. """ - def __init__(self, options: Mapping[str, Any], telemetry_messenger: ITelemetryMessenger): + def __init__(self, telemetry_messenger: ITelemetryMessenger): super(UsersPBA, self).__init__(telemetry_messenger, POST_BREACH_FILE_EXECUTION) self.filename = "" + def run(self, options: Dict) -> Iterable[PostBreachData]: + self._set_options(options) + return super().run(options) + + def _set_options(self, options: Dict): + # Required for attack telemetry + self.current_server = options["current_server"] + if is_windows_os(): # Add windows commands to PBA's if options["windows_filename"]: @@ -54,7 +62,7 @@ class UsersPBA(PBA): def _execute_default(self): if self.filename: - UsersPBA.download_pba_file(get_monkey_dir_path(), self.filename) + self.download_pba_file(get_monkey_dir_path(), self.filename) return super(UsersPBA, self)._execute_default() @staticmethod @@ -85,11 +93,11 @@ class UsersPBA(PBA): if not status: status = ScanStatus.USED - self._telemetry_messenger.send_telemetry( + self.telemetry_messenger.send_telemetry( T1105Telem( status, - WormConfiguration.current_server.split(":")[0], - get_interface_to_target(WormConfiguration.current_server.split(":")[0]), + self.current_server.split(":")[0], + get_interface_to_target(self.current_server.split(":")[0]), filename, ) ) diff --git a/monkey/monkey_island/cc/services/config.py b/monkey/monkey_island/cc/services/config.py index e636dfaab..863aa79a1 100644 --- a/monkey/monkey_island/cc/services/config.py +++ b/monkey/monkey_island/cc/services/config.py @@ -457,6 +457,8 @@ class ConfigService: "linux_filename": config.get(flat_linux_filename_field, ""), "windows_command": config.get(flat_windows_command_field, ""), "windows_filename": config.get(flat_windows_filename_field, ""), + # Current server is used for attack telemetry + "current_server": config.get("current_server"), } config["post_breach_actions"] = formatted_pbas_config diff --git a/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py b/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py index f8c51714b..4baa7f61d 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py +++ b/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py @@ -1,6 +1,6 @@ import logging import threading -from typing import Any, Dict, Iterable, List, Mapping, Sequence +from typing import Dict, Iterable, List, Sequence from infection_monkey.credential_collectors import LMHash, Password, SSHKeypair, Username from infection_monkey.i_puppet import ( @@ -57,9 +57,6 @@ class MockPuppet(IPuppet): else: return [PostBreachData(name, "pba command 2", ["pba result 2", False])] - def run_custom_pba(self, options: Mapping[str, Any]) -> PostBreachData: - pass - def ping(self, host: str, timeout: float = 1) -> PingScanData: logger.debug(f"run_ping({host}, {timeout})") if host == DOT_1: