From dbbdb508e3b2627272d935bd156f6a212738b029 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Tue, 29 Mar 2022 16:19:43 +0300 Subject: [PATCH] Agent: Change PBA constructor to accept telemetry messenger This change allows to run different PBA's with different telemetry messengers --- monkey/infection_monkey/monkey.py | 62 ++++++++++++++----- .../actions/change_file_privileges.py | 5 +- .../actions/clear_command_history.py | 9 +-- .../actions/collect_processes_list.py | 5 +- .../actions/communicate_as_backdoor_user.py | 5 +- .../post_breach/actions/discover_accounts.py | 8 ++- .../post_breach/actions/hide_files.py | 5 +- .../actions/modify_shell_startup_files.py | 5 +- .../post_breach/actions/schedule_jobs.py | 4 +- .../post_breach/actions/timestomping.py | 10 ++- .../post_breach/actions/use_signed_scripts.py | 9 ++- .../post_breach/actions/use_trap_command.py | 7 ++- .../post_breach/actions/users_custom_pba.py | 5 +- monkey/infection_monkey/post_breach/pba.py | 6 +- 14 files changed, 102 insertions(+), 43 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 478c8dde2..9576f76c0 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -248,26 +248,56 @@ class InfectionMonkey: ) puppet.load_plugin( - "CommunicateAsBackdoorUser", CommunicateAsBackdoorUser, PluginType.POST_BREACH_ACTION + "CommunicateAsBackdoorUser", + CommunicateAsBackdoorUser(self.telemetry_messenger), + PluginType.POST_BREACH_ACTION, ) puppet.load_plugin( - "ModifyShellStartupFiles", ModifyShellStartupFiles, PluginType.POST_BREACH_ACTION - ) - puppet.load_plugin("HiddenFiles", HiddenFiles, PluginType.POST_BREACH_ACTION) - puppet.load_plugin("TrapCommand", CommunicateAsBackdoorUser, PluginType.POST_BREACH_ACTION) - puppet.load_plugin("ChangeSetuidSetgid", ChangeSetuidSetgid, PluginType.POST_BREACH_ACTION) - puppet.load_plugin("ScheduleJobs", ScheduleJobs, PluginType.POST_BREACH_ACTION) - puppet.load_plugin("Timestomping", Timestomping, PluginType.POST_BREACH_ACTION) - puppet.load_plugin("AccountDiscovery", AccountDiscovery, PluginType.POST_BREACH_ACTION) - puppet.load_plugin( - "ProcessListCollection", ProcessListCollection, PluginType.POST_BREACH_ACTION - ) - puppet.load_plugin("TrapCommand", TrapCommand, PluginType.POST_BREACH_ACTION) - puppet.load_plugin( - "SignedScriptProxyExecution", SignedScriptProxyExecution, PluginType.POST_BREACH_ACTION + "ModifyShellStartupFiles", + ModifyShellStartupFiles(self.telemetry_messenger), + PluginType.POST_BREACH_ACTION, ) puppet.load_plugin( - "ClearCommandHistory", ClearCommandHistory, PluginType.POST_BREACH_ACTION + "HiddenFiles", HiddenFiles(self.telemetry_messenger), PluginType.POST_BREACH_ACTION + ) + puppet.load_plugin( + "TrapCommand", + CommunicateAsBackdoorUser(self.telemetry_messenger), + PluginType.POST_BREACH_ACTION, + ) + puppet.load_plugin( + "ChangeSetuidSetgid", + ChangeSetuidSetgid(self.telemetry_messenger), + PluginType.POST_BREACH_ACTION, + ) + puppet.load_plugin( + "ScheduleJobs", ScheduleJobs(self.telemetry_messenger), PluginType.POST_BREACH_ACTION + ) + puppet.load_plugin( + "Timestomping", Timestomping(self.telemetry_messenger), PluginType.POST_BREACH_ACTION + ) + puppet.load_plugin( + "AccountDiscovery", + AccountDiscovery(self.telemetry_messenger), + PluginType.POST_BREACH_ACTION, + ) + puppet.load_plugin( + "ProcessListCollection", + ProcessListCollection(self.telemetry_messenger), + PluginType.POST_BREACH_ACTION, + ) + puppet.load_plugin( + "TrapCommand", TrapCommand(self.telemetry_messenger), PluginType.POST_BREACH_ACTION + ) + puppet.load_plugin( + "SignedScriptProxyExecution", + SignedScriptProxyExecution(self.telemetry_messenger), + PluginType.POST_BREACH_ACTION, + ) + puppet.load_plugin( + "ClearCommandHistory", + ClearCommandHistory(self.telemetry_messenger), + PluginType.POST_BREACH_ACTION, ) puppet.load_plugin("ransomware", RansomwarePayload(), PluginType.PAYLOAD) diff --git a/monkey/infection_monkey/post_breach/actions/change_file_privileges.py b/monkey/infection_monkey/post_breach/actions/change_file_privileges.py index 87338e229..c560cc5d3 100644 --- a/monkey/infection_monkey/post_breach/actions/change_file_privileges.py +++ b/monkey/infection_monkey/post_breach/actions/change_file_privileges.py @@ -3,11 +3,12 @@ from infection_monkey.post_breach.pba import PBA from infection_monkey.post_breach.setuid_setgid.setuid_setgid import ( get_commands_to_change_setuid_setgid, ) +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger class ChangeSetuidSetgid(PBA): - def __init__(self): + def __init__(self, telemetry_messenger: ITelemetryMessenger): linux_cmds = get_commands_to_change_setuid_setgid() super(ChangeSetuidSetgid, self).__init__( - POST_BREACH_SETUID_SETGID, linux_cmd=" ".join(linux_cmds) + telemetry_messenger, POST_BREACH_SETUID_SETGID, linux_cmd=" ".join(linux_cmds) ) diff --git a/monkey/infection_monkey/post_breach/actions/clear_command_history.py b/monkey/infection_monkey/post_breach/actions/clear_command_history.py index 036c32d25..e6ab2d23e 100644 --- a/monkey/infection_monkey/post_breach/actions/clear_command_history.py +++ b/monkey/infection_monkey/post_breach/actions/clear_command_history.py @@ -6,20 +6,21 @@ from infection_monkey.post_breach.clear_command_history.clear_command_history im get_commands_to_clear_command_history, ) from infection_monkey.post_breach.pba import PBA +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger class ClearCommandHistory(PBA): - def __init__(self): - super().__init__(name=POST_BREACH_CLEAR_CMD_HISTORY) + def __init__(self, telemetry_messenger: ITelemetryMessenger): + super().__init__(telemetry_messenger, name=POST_BREACH_CLEAR_CMD_HISTORY) def run(self): - results = [pba.run() for pba in self.clear_command_history_PBA_list()] + results = [pba.run() for pba in self.clear_command_history_pba_list()] if results: # `self.command` is empty here self.pba_data.append(PostBreachData(self.name, self.command, results)) return self.pba_data - def clear_command_history_PBA_list(self): + def clear_command_history_pba_list(self): return self.CommandHistoryPBAGenerator().get_clear_command_history_pbas() class CommandHistoryPBAGenerator: diff --git a/monkey/infection_monkey/post_breach/actions/collect_processes_list.py b/monkey/infection_monkey/post_breach/actions/collect_processes_list.py index d0a5c5e0d..409583d18 100644 --- a/monkey/infection_monkey/post_breach/actions/collect_processes_list.py +++ b/monkey/infection_monkey/post_breach/actions/collect_processes_list.py @@ -5,6 +5,7 @@ import psutil from common.common_consts.post_breach_consts import POST_BREACH_PROCESS_LIST_COLLECTION from infection_monkey.i_puppet.i_puppet import PostBreachData from infection_monkey.post_breach.pba import PBA +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger logger = logging.getLogger(__name__) @@ -17,8 +18,8 @@ except NameError: class ProcessListCollection(PBA): - def __init__(self): - super().__init__(POST_BREACH_PROCESS_LIST_COLLECTION) + def __init__(self, telemetry_messenger: ITelemetryMessenger): + super().__init__(telemetry_messenger, POST_BREACH_PROCESS_LIST_COLLECTION) def run(self): """ diff --git a/monkey/infection_monkey/post_breach/actions/communicate_as_backdoor_user.py b/monkey/infection_monkey/post_breach/actions/communicate_as_backdoor_user.py index e4523f0fd..60990d67a 100644 --- a/monkey/infection_monkey/post_breach/actions/communicate_as_backdoor_user.py +++ b/monkey/infection_monkey/post_breach/actions/communicate_as_backdoor_user.py @@ -7,6 +7,7 @@ import subprocess from common.common_consts.post_breach_consts import POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER from infection_monkey.i_puppet.i_puppet import PostBreachData from infection_monkey.post_breach.pba import PBA +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger from infection_monkey.utils.auto_new_user_factory import create_auto_new_user from infection_monkey.utils.environment import is_windows_os from infection_monkey.utils.new_user_error import NewUserError @@ -33,9 +34,9 @@ class CommunicateAsBackdoorUser(PBA): are created. """ - def __init__(self): + def __init__(self, telemetry_messenger: ITelemetryMessenger): super(CommunicateAsBackdoorUser, self).__init__( - name=POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER + telemetry_messenger, name=POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER ) def run(self): diff --git a/monkey/infection_monkey/post_breach/actions/discover_accounts.py b/monkey/infection_monkey/post_breach/actions/discover_accounts.py index 8fdebd0df..a153cf5b6 100644 --- a/monkey/infection_monkey/post_breach/actions/discover_accounts.py +++ b/monkey/infection_monkey/post_breach/actions/discover_accounts.py @@ -3,11 +3,15 @@ from infection_monkey.post_breach.account_discovery.account_discovery import ( get_commands_to_discover_accounts, ) from infection_monkey.post_breach.pba import PBA +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger class AccountDiscovery(PBA): - def __init__(self): + def __init__(self, telemetry_messenger: ITelemetryMessenger): linux_cmds, windows_cmds = get_commands_to_discover_accounts() super().__init__( - POST_BREACH_ACCOUNT_DISCOVERY, linux_cmd=" ".join(linux_cmds), windows_cmd=windows_cmds + telemetry_messenger, + POST_BREACH_ACCOUNT_DISCOVERY, + linux_cmd=" ".join(linux_cmds), + windows_cmd=windows_cmds, ) diff --git a/monkey/infection_monkey/post_breach/actions/hide_files.py b/monkey/infection_monkey/post_breach/actions/hide_files.py index e3123192c..457b9dafe 100644 --- a/monkey/infection_monkey/post_breach/actions/hide_files.py +++ b/monkey/infection_monkey/post_breach/actions/hide_files.py @@ -1,6 +1,7 @@ from common.common_consts.post_breach_consts import POST_BREACH_HIDDEN_FILES from infection_monkey.i_puppet.i_puppet import PostBreachData from infection_monkey.post_breach.pba import PBA +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger from infection_monkey.utils.environment import is_windows_os from infection_monkey.utils.hidden_files import ( cleanup_hidden_files, @@ -17,8 +18,8 @@ class HiddenFiles(PBA): This PBA attempts to create hidden files and folders. """ - def __init__(self): - super(HiddenFiles, self).__init__(name=POST_BREACH_HIDDEN_FILES) + def __init__(self, telemetry_messenger: ITelemetryMessenger): + super(HiddenFiles, self).__init__(telemetry_messenger, name=POST_BREACH_HIDDEN_FILES) def run(self): # create hidden files and folders diff --git a/monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py b/monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py index 5d3c3c5ea..4d755567b 100644 --- a/monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py +++ b/monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py @@ -6,6 +6,7 @@ from infection_monkey.post_breach.pba import PBA from infection_monkey.post_breach.shell_startup_files.shell_startup_files_modification import ( get_commands_to_modify_shell_startup_files, ) +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger class ModifyShellStartupFiles(PBA): @@ -15,8 +16,8 @@ class ModifyShellStartupFiles(PBA): and profile.ps1 in windows. """ - def __init__(self): - super().__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION) + def __init__(self, telemetry_messenger: ITelemetryMessenger): + super().__init__(telemetry_messenger, name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION) def run(self): results = [pba.run() for pba in self.modify_shell_startup_PBA_list()] diff --git a/monkey/infection_monkey/post_breach/actions/schedule_jobs.py b/monkey/infection_monkey/post_breach/actions/schedule_jobs.py index 37649488b..8aeb0b42d 100644 --- a/monkey/infection_monkey/post_breach/actions/schedule_jobs.py +++ b/monkey/infection_monkey/post_breach/actions/schedule_jobs.py @@ -4,6 +4,7 @@ from infection_monkey.post_breach.job_scheduling.job_scheduling import ( remove_scheduled_jobs, ) from infection_monkey.post_breach.pba import PBA +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger class ScheduleJobs(PBA): @@ -11,10 +12,11 @@ class ScheduleJobs(PBA): This PBA attempts to schedule jobs on the system. """ - def __init__(self): + def __init__(self, telemetry_messenger: ITelemetryMessenger): linux_cmds, windows_cmds = get_commands_to_schedule_jobs() super(ScheduleJobs, self).__init__( + telemetry_messenger, name=POST_BREACH_JOB_SCHEDULING, linux_cmd=" ".join(linux_cmds), windows_cmd=windows_cmds, diff --git a/monkey/infection_monkey/post_breach/actions/timestomping.py b/monkey/infection_monkey/post_breach/actions/timestomping.py index ece987107..3e7c61f59 100644 --- a/monkey/infection_monkey/post_breach/actions/timestomping.py +++ b/monkey/infection_monkey/post_breach/actions/timestomping.py @@ -1,9 +1,15 @@ from common.common_consts.post_breach_consts import POST_BREACH_TIMESTOMPING from infection_monkey.post_breach.pba import PBA from infection_monkey.post_breach.timestomping.timestomping import get_timestomping_commands +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger class Timestomping(PBA): - def __init__(self): + def __init__(self, telemetry_messenger: ITelemetryMessenger): linux_cmds, windows_cmds = get_timestomping_commands() - super().__init__(POST_BREACH_TIMESTOMPING, linux_cmd=linux_cmds, windows_cmd=windows_cmds) + super().__init__( + telemetry_messenger, + POST_BREACH_TIMESTOMPING, + linux_cmd=linux_cmds, + windows_cmd=windows_cmds, + ) diff --git a/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py b/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py index 75ede03ee..d7323b54e 100644 --- a/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py +++ b/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py @@ -7,15 +7,20 @@ from infection_monkey.post_breach.signed_script_proxy.signed_script_proxy import cleanup_changes, get_commands_to_proxy_execution_using_signed_script, ) +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger from infection_monkey.utils.environment import is_windows_os logger = logging.getLogger(__name__) class SignedScriptProxyExecution(PBA): - def __init__(self): + def __init__(self, telemetry_messenger: ITelemetryMessenger): windows_cmds = get_commands_to_proxy_execution_using_signed_script() - super().__init__(POST_BREACH_SIGNED_SCRIPT_PROXY_EXEC, windows_cmd=" ".join(windows_cmds)) + super().__init__( + telemetry_messenger, + POST_BREACH_SIGNED_SCRIPT_PROXY_EXEC, + windows_cmd=" ".join(windows_cmds), + ) def run(self): original_comspec = "" diff --git a/monkey/infection_monkey/post_breach/actions/use_trap_command.py b/monkey/infection_monkey/post_breach/actions/use_trap_command.py index 879db77bf..8dfbc9f5e 100644 --- a/monkey/infection_monkey/post_breach/actions/use_trap_command.py +++ b/monkey/infection_monkey/post_breach/actions/use_trap_command.py @@ -1,9 +1,12 @@ from common.common_consts.post_breach_consts import POST_BREACH_TRAP_COMMAND from infection_monkey.post_breach.pba import PBA from infection_monkey.post_breach.trap_command.trap_command import get_trap_commands +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger class TrapCommand(PBA): - def __init__(self): + def __init__(self, telemetry_messenger: ITelemetryMessenger): linux_cmds = get_trap_commands() - super(TrapCommand, self).__init__(POST_BREACH_TRAP_COMMAND, linux_cmd=" ".join(linux_cmds)) + super(TrapCommand, self).__init__( + telemetry_messenger, POST_BREACH_TRAP_COMMAND, linux_cmd=" ".join(linux_cmds) + ) diff --git a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py b/monkey/infection_monkey/post_breach/actions/users_custom_pba.py index b1ccec85c..91475e66d 100644 --- a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py +++ b/monkey/infection_monkey/post_breach/actions/users_custom_pba.py @@ -8,6 +8,7 @@ from infection_monkey.control import ControlClient 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 +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger from infection_monkey.utils.environment import is_windows_os from infection_monkey.utils.monkey_dir import get_monkey_dir_path @@ -23,8 +24,8 @@ class UsersPBA(PBA): Defines user's configured post breach action. """ - def __init__(self): - super(UsersPBA, self).__init__(POST_BREACH_FILE_EXECUTION) + def __init__(self, telemetry_messenger: ITelemetryMessenger): + super(UsersPBA, self).__init__(telemetry_messenger, POST_BREACH_FILE_EXECUTION) self.filename = "" if not is_windows_os(): diff --git a/monkey/infection_monkey/post_breach/pba.py b/monkey/infection_monkey/post_breach/pba.py index 769ff2de0..9222d5e1a 100644 --- a/monkey/infection_monkey/post_breach/pba.py +++ b/monkey/infection_monkey/post_breach/pba.py @@ -6,7 +6,6 @@ from common.utils.attack_utils import ScanStatus from infection_monkey.i_puppet.i_puppet import PostBreachData from infection_monkey.telemetry.attack.t1064_telem import T1064Telem from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger -from infection_monkey.telemetry.post_breach_telem import PostBreachTelem from infection_monkey.utils.environment import is_windows_os logger = logging.getLogger(__name__) @@ -18,7 +17,9 @@ class PBA: machine. """ - def __init__(self, name="unknown", linux_cmd="", windows_cmd=""): + def __init__( + self, telemetry_messenger: ITelemetryMessenger, name="unknown", linux_cmd="", windows_cmd="" + ): """ :param name: Name of post breach action. :param linux_cmd: Command that will be executed on breached machine @@ -27,6 +28,7 @@ class PBA: self.command = PBA.choose_command(linux_cmd, windows_cmd) self.name = name self.pba_data = [] + self.telemetry_messenger = telemetry_messenger def run(self) -> Iterable[PostBreachData]: """