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"