From 23b8c351fbe020ec1ff9f26333a968c9cc0255b3 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Fri, 25 Mar 2022 16:38:13 +0200 Subject: [PATCH 1/8] Island, Agent: Add custom user PBA to puppet and master --- monkey/infection_monkey/i_puppet/i_puppet.py | 10 +++- .../master/automated_master.py | 16 ++++-- .../post_breach/custom_pba/__init__.py | 0 .../users_custom_pba.py | 49 ++++++++++--------- monkey/monkey_island/cc/services/config.py | 2 +- .../infection_monkey/master/mock_puppet.py | 5 +- .../actions/test_users_custom_pba.py | 2 +- 7 files changed, 53 insertions(+), 31 deletions(-) create mode 100644 monkey/infection_monkey/post_breach/custom_pba/__init__.py rename monkey/infection_monkey/post_breach/{actions => custom_pba}/users_custom_pba.py (69%) diff --git a/monkey/infection_monkey/i_puppet/i_puppet.py b/monkey/infection_monkey/i_puppet/i_puppet.py index c4f46d792..fd42d5c58 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 Dict, Iterable, List, Mapping, Sequence +from typing import Any, Dict, Iterable, List, Mapping, Sequence from infection_monkey.model import VictimHost @@ -67,6 +67,14 @@ 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 e5269fa64..42086fee4 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -1,9 +1,10 @@ import logging import threading import time -from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple +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 @@ -154,9 +155,9 @@ class AutomatedMaster(IMaster): ), ) pba_thread = create_daemon_thread( - target=self._run_plugins, + target=self._run_PBAs, name="PBAThread", - args=(config["post_breach_actions"].items(), "post-breach action", self._run_pba), + args=(config["post_breach_actions"].items(), self._run_pba, config["custom_pbas"]), ) credential_collector_thread.start() @@ -212,6 +213,15 @@ class AutomatedMaster(IMaster): self._puppet.run_payload(name, options, self._stop) + 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) + def _run_plugins( self, plugins: Iterable[Any], plugin_type: str, callback: Callable[[Any], None] ): diff --git a/monkey/infection_monkey/post_breach/custom_pba/__init__.py b/monkey/infection_monkey/post_breach/custom_pba/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py b/monkey/infection_monkey/post_breach/custom_pba/users_custom_pba.py similarity index 69% rename from monkey/infection_monkey/post_breach/actions/users_custom_pba.py rename to monkey/infection_monkey/post_breach/custom_pba/users_custom_pba.py index 91475e66d..81b6e9ab2 100644 --- a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py +++ b/monkey/infection_monkey/post_breach/custom_pba/users_custom_pba.py @@ -1,5 +1,6 @@ import logging import os +from typing import Any, Mapping from common.common_consts.post_breach_consts import POST_BREACH_FILE_EXECUTION from common.utils.attack_utils import ScanStatus @@ -24,32 +25,32 @@ class UsersPBA(PBA): Defines user's configured post breach action. """ - def __init__(self, telemetry_messenger: ITelemetryMessenger): + def __init__(self, options: Mapping[str, Any], telemetry_messenger: ITelemetryMessenger): super(UsersPBA, self).__init__(telemetry_messenger, POST_BREACH_FILE_EXECUTION) self.filename = "" - if not is_windows_os(): - # Add linux commands to PBA's - if WormConfiguration.PBA_linux_filename: - self.filename = WormConfiguration.PBA_linux_filename - if WormConfiguration.custom_PBA_linux_cmd: - # Add change dir command, because user will try to access his file - self.command = ( - DIR_CHANGE_LINUX % get_monkey_dir_path() - ) + WormConfiguration.custom_PBA_linux_cmd - elif WormConfiguration.custom_PBA_linux_cmd: - self.command = WormConfiguration.custom_PBA_linux_cmd - else: + if is_windows_os(): # Add windows commands to PBA's - if WormConfiguration.PBA_windows_filename: - self.filename = WormConfiguration.PBA_windows_filename - if WormConfiguration.custom_PBA_windows_cmd: + if options["windows_filename"]: + self.filename = options["windows_filename"] + if options["windows_command"]: # Add change dir command, because user will try to access his file - self.command = ( - DIR_CHANGE_WINDOWS % get_monkey_dir_path() - ) + WormConfiguration.custom_PBA_windows_cmd - elif WormConfiguration.custom_PBA_windows_cmd: - self.command = WormConfiguration.custom_PBA_windows_cmd + self.command = (DIR_CHANGE_WINDOWS % get_monkey_dir_path()) + options[ + "windows_command" + ] + elif options["windows_command"]: + self.command = options["windows_command"] + else: + # Add linux commands to PBA's + if options["linux_filename"]: + self.filename = options["linux_filename"] + if options["linux_command"]: + # Add change dir command, because user will try to access his file + self.command = (DIR_CHANGE_LINUX % get_monkey_dir_path()) + options[ + "linux_command" + ] + elif options["linux_command"]: + self.command = options["linux_command"] def _execute_default(self): if self.filename: @@ -57,12 +58,12 @@ class UsersPBA(PBA): return super(UsersPBA, self)._execute_default() @staticmethod - def should_run(class_name): + def should_run(options): if not is_windows_os(): - if WormConfiguration.PBA_linux_filename or WormConfiguration.custom_PBA_linux_cmd: + if options["linux_filename"] or options["linux_command"]: return True else: - if WormConfiguration.PBA_windows_filename or WormConfiguration.custom_PBA_windows_cmd: + if options["windows_filename"] or options["windows_command"]: return True return False diff --git a/monkey/monkey_island/cc/services/config.py b/monkey/monkey_island/cc/services/config.py index c5f78e62d..e636dfaab 100644 --- a/monkey/monkey_island/cc/services/config.py +++ b/monkey/monkey_island/cc/services/config.py @@ -452,7 +452,7 @@ class ConfigService: for pba in config.get("post_breach_actions", []): formatted_pbas_config[pba] = {} - formatted_pbas_config["Custom"] = { + config["custom_pbas"] = { "linux_command": config.get(flat_linux_command_field, ""), "linux_filename": config.get(flat_linux_filename_field, ""), "windows_command": config.get(flat_windows_command_field, ""), 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 4baa7f61d..f8c51714b 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 Dict, Iterable, List, Sequence +from typing import Any, Dict, Iterable, List, Mapping, Sequence from infection_monkey.credential_collectors import LMHash, Password, SSHKeypair, Username from infection_monkey.i_puppet import ( @@ -57,6 +57,9 @@ 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: diff --git a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py index 2618f8d93..e8cdd0333 100644 --- a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py +++ b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py @@ -2,7 +2,7 @@ from unittest.mock import MagicMock import pytest -from infection_monkey.post_breach.actions.users_custom_pba import UsersPBA +from infection_monkey.post_breach.custom_pba.users_custom_pba import UsersPBA MONKEY_DIR_PATH = "/dir/to/monkey/" CUSTOM_LINUX_CMD = "command-for-linux" From 24915ba797f27d48fa54b2ce46c510fdd5df40c3 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 30 Mar 2022 14:13:51 +0300 Subject: [PATCH 2/8] 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: From 079d768f73c51be04446b1e1f0429b3198116bad Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 30 Mar 2022 14:17:31 +0300 Subject: [PATCH 3/8] Agent: Rename UsersPBA to CustomPBA for consistency --- monkey/infection_monkey/monkey.py | 4 ++-- .../custom_pba/{users_custom_pba.py => custom_pba.py} | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) rename monkey/infection_monkey/post_breach/custom_pba/{users_custom_pba.py => custom_pba.py} (96%) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 1863da03f..b500fa27f 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -51,7 +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.post_breach.custom_pba.custom_pba import CustomPBA from infection_monkey.puppet.puppet import Puppet from infection_monkey.system_singleton import SystemSingleton from infection_monkey.telemetry.attack.t1106_telem import T1106Telem @@ -317,7 +317,7 @@ class InfectionMonkey: PluginType.POST_BREACH_ACTION, ) puppet.load_plugin( - "CustomPBA", UsersPBA(self._telemetry_messenger), PluginType.POST_BREACH_ACTION + "CustomPBA", CustomPBA(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/custom_pba.py similarity index 96% rename from monkey/infection_monkey/post_breach/custom_pba/users_custom_pba.py rename to monkey/infection_monkey/post_breach/custom_pba/custom_pba.py index 5f9100d98..51e9ef7a1 100644 --- a/monkey/infection_monkey/post_breach/custom_pba/users_custom_pba.py +++ b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py @@ -20,13 +20,13 @@ DIR_CHANGE_WINDOWS = "cd %s & " DIR_CHANGE_LINUX = "cd %s ; " -class UsersPBA(PBA): +class CustomPBA(PBA): """ Defines user's configured post breach action. """ def __init__(self, telemetry_messenger: ITelemetryMessenger): - super(UsersPBA, self).__init__(telemetry_messenger, POST_BREACH_FILE_EXECUTION) + super(CustomPBA, self).__init__(telemetry_messenger, POST_BREACH_FILE_EXECUTION) self.filename = "" def run(self, options: Dict) -> Iterable[PostBreachData]: @@ -63,7 +63,7 @@ class UsersPBA(PBA): def _execute_default(self): if self.filename: self.download_pba_file(get_monkey_dir_path(), self.filename) - return super(UsersPBA, self)._execute_default() + return super(CustomPBA, self)._execute_default() @staticmethod def should_run(options): From 67543ef91a432079abe36565e327342169268d75 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 30 Mar 2022 14:26:45 +0300 Subject: [PATCH 4/8] Agent: Add a custom PBA run check We only want to run the custom PBA if commands are specified --- monkey/infection_monkey/master/automated_master.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index c76a8938c..3a4ef6835 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -15,6 +15,7 @@ from infection_monkey.telemetry.post_breach_telem import PostBreachTelem from infection_monkey.utils.threading import create_daemon_thread, interruptible_iter from infection_monkey.utils.timer import Timer +from ..post_breach.custom_pba.custom_pba import CustomPBA from . import Exploiter, IPScanner, Propagator CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC = 5 @@ -213,7 +214,8 @@ class AutomatedMaster(IMaster): ): self._run_plugins(plugins, "post-breach action", callback) - self._run_plugins([("CustomPBA", custom_pba_options)], "post-breach action", callback) + if CustomPBA.should_run(custom_pba_options): + 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] From 1f31e96adbb6bbb0898ed8ee1cb2850b6591a6db Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 30 Mar 2022 15:48:16 +0300 Subject: [PATCH 5/8] Agent: Make custom PBA related imports shorter --- monkey/infection_monkey/master/automated_master.py | 2 +- monkey/infection_monkey/monkey.py | 2 +- monkey/infection_monkey/post_breach/custom_pba/__init__.py | 1 + 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 3a4ef6835..88841f401 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -9,13 +9,13 @@ from infection_monkey.i_master import IMaster from infection_monkey.i_puppet import IPuppet from infection_monkey.model import VictimHostFactory from infection_monkey.network import NetworkInterface +from infection_monkey.post_breach.custom_pba import CustomPBA from infection_monkey.telemetry.credentials_telem import CredentialsTelem from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger from infection_monkey.telemetry.post_breach_telem import PostBreachTelem from infection_monkey.utils.threading import create_daemon_thread, interruptible_iter from infection_monkey.utils.timer import Timer -from ..post_breach.custom_pba.custom_pba import CustomPBA from . import Exploiter, IPScanner, Propagator CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC = 5 diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index b500fa27f..96f1873bc 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -51,7 +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.custom_pba import CustomPBA +from infection_monkey.post_breach.custom_pba import CustomPBA from infection_monkey.puppet.puppet import Puppet from infection_monkey.system_singleton import SystemSingleton from infection_monkey.telemetry.attack.t1106_telem import T1106Telem diff --git a/monkey/infection_monkey/post_breach/custom_pba/__init__.py b/monkey/infection_monkey/post_breach/custom_pba/__init__.py index e69de29bb..723810fc0 100644 --- a/monkey/infection_monkey/post_breach/custom_pba/__init__.py +++ b/monkey/infection_monkey/post_breach/custom_pba/__init__.py @@ -0,0 +1 @@ +from .custom_pba import CustomPBA From 2e3a718469111dc971f2e1edeb17c574d669f4d7 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 30 Mar 2022 15:48:43 +0300 Subject: [PATCH 6/8] Agent: Fix custom PBA related unit tests --- .../actions/test_users_custom_pba.py | 149 ++++++++---------- .../monkey_island/cc/services/test_config.py | 19 ++- 2 files changed, 81 insertions(+), 87 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py index e8cdd0333..598f3412b 100644 --- a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py +++ b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py @@ -2,19 +2,20 @@ from unittest.mock import MagicMock import pytest -from infection_monkey.post_breach.custom_pba.users_custom_pba import UsersPBA +from infection_monkey.post_breach.custom_pba.custom_pba import CustomPBA MONKEY_DIR_PATH = "/dir/to/monkey/" CUSTOM_LINUX_CMD = "command-for-linux" CUSTOM_LINUX_FILENAME = "filename-for-linux" CUSTOM_WINDOWS_CMD = "command-for-windows" CUSTOM_WINDOWS_FILENAME = "filename-for-windows" +CUSTOM_SERVER = "10.10.10.10:5000" -@pytest.fixture +@pytest.fixture(autouse=True) def fake_monkey_dir_path(monkeypatch): monkeypatch.setattr( - "infection_monkey.post_breach.actions.users_custom_pba.get_monkey_dir_path", + "infection_monkey.post_breach.custom_pba.custom_pba.get_monkey_dir_path", lambda: MONKEY_DIR_PATH, ) @@ -22,7 +23,7 @@ def fake_monkey_dir_path(monkeypatch): @pytest.fixture def set_os_linux(monkeypatch): monkeypatch.setattr( - "infection_monkey.post_breach.actions.users_custom_pba.is_windows_os", + "infection_monkey.post_breach.custom_pba.custom_pba.is_windows_os", lambda: False, ) @@ -30,106 +31,92 @@ def set_os_linux(monkeypatch): @pytest.fixture def set_os_windows(monkeypatch): monkeypatch.setattr( - "infection_monkey.post_breach.actions.users_custom_pba.is_windows_os", + "infection_monkey.post_breach.custom_pba.custom_pba.is_windows_os", lambda: True, ) @pytest.fixture -def mock_UsersPBA_linux_custom_file_and_cmd(set_os_linux, fake_monkey_dir_path, monkeypatch): - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", - CUSTOM_LINUX_CMD, - ) - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.PBA_linux_filename", - CUSTOM_LINUX_FILENAME, - ) - return UsersPBA(MagicMock()) +def fake_custom_pba_linux_options(): + return { + "linux_command": CUSTOM_LINUX_CMD, + "linux_filename": CUSTOM_LINUX_FILENAME, + "windows_command": "", + "windows_filename": "", + # Current server is used for attack telemetry + "current_server": CUSTOM_SERVER, + } -def test_command_linux_custom_file_and_cmd( - mock_UsersPBA_linux_custom_file_and_cmd, -): +def test_command_linux_custom_file_and_cmd(fake_custom_pba_linux_options, set_os_linux): + pba = CustomPBA(MagicMock()) + pba._set_options(fake_custom_pba_linux_options) expected_command = f"cd {MONKEY_DIR_PATH} ; {CUSTOM_LINUX_CMD}" - assert mock_UsersPBA_linux_custom_file_and_cmd.command == expected_command + assert pba.command == expected_command + assert pba.filename == CUSTOM_LINUX_FILENAME @pytest.fixture -def mock_UsersPBA_windows_custom_file_and_cmd(set_os_windows, fake_monkey_dir_path, monkeypatch): - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", - CUSTOM_WINDOWS_CMD, - ) - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.PBA_windows_filename", - CUSTOM_WINDOWS_FILENAME, - ) - return UsersPBA(MagicMock()) +def fake_custom_pba_windows_options(): + return { + "linux_command": "", + "linux_filename": "", + "windows_command": CUSTOM_WINDOWS_CMD, + "windows_filename": CUSTOM_WINDOWS_FILENAME, + # Current server is used for attack telemetry + "current_server": CUSTOM_SERVER, + } -def test_command_windows_custom_file_and_cmd( - mock_UsersPBA_windows_custom_file_and_cmd, -): +def test_command_windows_custom_file_and_cmd(fake_custom_pba_windows_options, set_os_windows): + + pba = CustomPBA(MagicMock()) + pba._set_options(fake_custom_pba_windows_options) expected_command = f"cd {MONKEY_DIR_PATH} & {CUSTOM_WINDOWS_CMD}" - assert mock_UsersPBA_windows_custom_file_and_cmd.command == expected_command + assert pba.command == expected_command + assert pba.filename == CUSTOM_WINDOWS_FILENAME @pytest.fixture -def mock_UsersPBA_linux_custom_file(set_os_linux, fake_monkey_dir_path, monkeypatch): - monkeypatch.setattr("infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", None) - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.PBA_linux_filename", - CUSTOM_LINUX_FILENAME, - ) - return UsersPBA(MagicMock()) +def fake_options_files_only(): + return { + "linux_command": "", + "linux_filename": CUSTOM_LINUX_FILENAME, + "windows_command": "", + "windows_filename": CUSTOM_WINDOWS_FILENAME, + # Current server is used for attack telemetry + "current_server": CUSTOM_SERVER, + } -def test_command_linux_custom_file(mock_UsersPBA_linux_custom_file): - expected_command = "" - assert mock_UsersPBA_linux_custom_file.command == expected_command +@pytest.mark.parametrize("os", [set_os_linux, set_os_windows]) +def test_files_only(fake_options_files_only, os): + pba = CustomPBA(MagicMock()) + pba._set_options(fake_options_files_only) + assert pba.command == "" @pytest.fixture -def mock_UsersPBA_windows_custom_file(set_os_windows, fake_monkey_dir_path, monkeypatch): - monkeypatch.setattr("infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", None) - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.PBA_windows_filename", - CUSTOM_WINDOWS_FILENAME, - ) - return UsersPBA(MagicMock()) +def fake_options_commands_only(): + return { + "linux_command": CUSTOM_LINUX_CMD, + "linux_filename": "", + "windows_command": CUSTOM_WINDOWS_CMD, + "windows_filename": "", + # Current server is used for attack telemetry + "current_server": CUSTOM_SERVER, + } -def test_command_windows_custom_file(mock_UsersPBA_windows_custom_file): - expected_command = "" - assert mock_UsersPBA_windows_custom_file.command == expected_command +def test_commands_only(fake_options_commands_only, set_os_linux): + pba = CustomPBA(MagicMock()) + pba._set_options(fake_options_commands_only) + assert pba.command == CUSTOM_LINUX_CMD + assert pba.filename == "" -@pytest.fixture -def mock_UsersPBA_linux_custom_cmd(set_os_linux, fake_monkey_dir_path, monkeypatch): - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", - CUSTOM_LINUX_CMD, - ) - monkeypatch.setattr("infection_monkey.config.WormConfiguration.PBA_linux_filename", None) - return UsersPBA(MagicMock()) - - -def test_command_linux_custom_cmd(mock_UsersPBA_linux_custom_cmd): - expected_command = CUSTOM_LINUX_CMD - assert mock_UsersPBA_linux_custom_cmd.command == expected_command - - -@pytest.fixture -def mock_UsersPBA_windows_custom_cmd(set_os_windows, fake_monkey_dir_path, monkeypatch): - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", - CUSTOM_WINDOWS_CMD, - ) - monkeypatch.setattr("infection_monkey.config.WormConfiguration.PBA_windows_filename", None) - return UsersPBA(MagicMock()) - - -def test_command_windows_custom_cmd(mock_UsersPBA_windows_custom_cmd): - expected_command = CUSTOM_WINDOWS_CMD - assert mock_UsersPBA_windows_custom_cmd.command == expected_command +def test_commands_only_windows(fake_options_commands_only, set_os_windows): + pba = CustomPBA(MagicMock()) + pba._set_options(fake_options_commands_only) + assert pba.command == CUSTOM_WINDOWS_CMD + assert pba.filename == "" 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 ae0a44cdc..b49007eb0 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 @@ -66,12 +66,6 @@ def test_format_config_for_agent__pbas(flat_monkey_config): "ScheduleJobs": {}, "Timestomping": {}, "AccountDiscovery": {}, - "Custom": { - "linux_command": "bash test.sh", - "windows_command": "powershell test.ps1", - "linux_filename": "test.sh", - "windows_filename": "test.ps1", - }, } ConfigService.format_flat_config_for_agent(flat_monkey_config) @@ -84,6 +78,19 @@ def test_format_config_for_agent__pbas(flat_monkey_config): assert "PBA_windows_filename" not in flat_monkey_config +def test_format_config_for_custom_pbas(flat_monkey_config): + custom_config = { + "linux_command": "bash test.sh", + "windows_command": "powershell test.ps1", + "linux_filename": "test.sh", + "windows_filename": "test.ps1", + "current_server": "10.197.94.72:5000", + } + ConfigService.format_flat_config_for_agent(flat_monkey_config) + + assert flat_monkey_config["custom_pbas"] == custom_config + + def test_get_config_propagation_credentials_from_flat_config(flat_monkey_config): expected_creds = { "exploit_lm_hash_list": ["lm_hash_1", "lm_hash_2"], From e855d2ed34758b216246ebda82db2b882d5abb84 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 30 Mar 2022 16:07:14 +0300 Subject: [PATCH 7/8] Agent: Remove unused pba properties in config.py --- monkey/infection_monkey/config.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/monkey/infection_monkey/config.py b/monkey/infection_monkey/config.py index 8e9ffce8f..7b8d793cf 100644 --- a/monkey/infection_monkey/config.py +++ b/monkey/infection_monkey/config.py @@ -89,15 +89,6 @@ class Configuration(object): keep_tunnel_open_time = 60 - ########################### - # post breach actions - ########################### - post_breach_actions = [] - custom_PBA_linux_cmd = "" - custom_PBA_windows_cmd = "" - PBA_linux_filename = None - PBA_windows_filename = None - ########################### # testing configuration ########################### From 29a545a58fc7d845fbbd12e28cb80ba1433559c6 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 30 Mar 2022 16:37:19 +0300 Subject: [PATCH 8/8] Agent: Move the decision if custom pba should run to master --- monkey/infection_monkey/master/automated_master.py | 4 ++-- monkey/infection_monkey/master/option_parsing.py | 13 +++++++++++++ .../post_breach/custom_pba/custom_pba.py | 10 ---------- 3 files changed, 15 insertions(+), 12 deletions(-) create mode 100644 monkey/infection_monkey/master/option_parsing.py diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 88841f401..0ef31129c 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -9,7 +9,6 @@ from infection_monkey.i_master import IMaster from infection_monkey.i_puppet import IPuppet from infection_monkey.model import VictimHostFactory from infection_monkey.network import NetworkInterface -from infection_monkey.post_breach.custom_pba import CustomPBA from infection_monkey.telemetry.credentials_telem import CredentialsTelem from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger from infection_monkey.telemetry.post_breach_telem import PostBreachTelem @@ -17,6 +16,7 @@ from infection_monkey.utils.threading import create_daemon_thread, interruptible from infection_monkey.utils.timer import Timer from . import Exploiter, IPScanner, Propagator +from .option_parsing import custom_pba_is_enabled CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC = 5 CHECK_FOR_TERMINATE_INTERVAL_SEC = CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC / 5 @@ -214,7 +214,7 @@ class AutomatedMaster(IMaster): ): self._run_plugins(plugins, "post-breach action", callback) - if CustomPBA.should_run(custom_pba_options): + if custom_pba_is_enabled(custom_pba_options): self._run_plugins([("CustomPBA", custom_pba_options)], "post-breach action", callback) def _run_plugins( diff --git a/monkey/infection_monkey/master/option_parsing.py b/monkey/infection_monkey/master/option_parsing.py new file mode 100644 index 000000000..c35bf6303 --- /dev/null +++ b/monkey/infection_monkey/master/option_parsing.py @@ -0,0 +1,13 @@ +from typing import Dict + +from infection_monkey.utils.environment import is_windows_os + + +def custom_pba_is_enabled(pba_options: Dict) -> bool: + if not is_windows_os(): + if pba_options["linux_command"]: + return True + else: + if pba_options["windows_command"]: + return True + return False diff --git a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py index 51e9ef7a1..453dfb6ed 100644 --- a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py +++ b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py @@ -65,16 +65,6 @@ class CustomPBA(PBA): self.download_pba_file(get_monkey_dir_path(), self.filename) return super(CustomPBA, self)._execute_default() - @staticmethod - def should_run(options): - if not is_windows_os(): - if options["linux_filename"] or options["linux_command"]: - return True - else: - if options["windows_filename"] or options["windows_command"]: - return True - return False - def download_pba_file(self, dst_dir, filename): """ Handles post breach action file download