From 40b1ae005846fe620fd3138c8832be42580749e2 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 30 Mar 2022 13:37:47 +0530 Subject: [PATCH] Agent: Modify puppet to run PBAs instead of using the mock puppet --- .../post_breach/actions/clear_command_history.py | 3 ++- .../post_breach/actions/collect_processes_list.py | 3 ++- .../post_breach/actions/communicate_as_backdoor_user.py | 3 ++- monkey/infection_monkey/post_breach/actions/hide_files.py | 6 ++++-- .../post_breach/actions/modify_shell_startup_files.py | 3 ++- .../infection_monkey/post_breach/actions/schedule_jobs.py | 6 ++++-- .../post_breach/actions/use_signed_scripts.py | 5 +++-- monkey/infection_monkey/puppet/puppet.py | 3 ++- 8 files changed, 21 insertions(+), 11 deletions(-) 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 e6ab2d23e..3ef363121 100644 --- a/monkey/infection_monkey/post_breach/actions/clear_command_history.py +++ b/monkey/infection_monkey/post_breach/actions/clear_command_history.py @@ -1,4 +1,5 @@ import subprocess +from typing import Dict from common.common_consts.post_breach_consts import POST_BREACH_CLEAR_CMD_HISTORY from infection_monkey.i_puppet.i_puppet import PostBreachData @@ -13,7 +14,7 @@ class ClearCommandHistory(PBA): def __init__(self, telemetry_messenger: ITelemetryMessenger): super().__init__(telemetry_messenger, name=POST_BREACH_CLEAR_CMD_HISTORY) - def run(self): + def run(self, options: Dict): results = [pba.run() for pba in self.clear_command_history_pba_list()] if results: # `self.command` is empty here 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 409583d18..78102c595 100644 --- a/monkey/infection_monkey/post_breach/actions/collect_processes_list.py +++ b/monkey/infection_monkey/post_breach/actions/collect_processes_list.py @@ -1,4 +1,5 @@ import logging +from typing import Dict import psutil @@ -21,7 +22,7 @@ class ProcessListCollection(PBA): def __init__(self, telemetry_messenger: ITelemetryMessenger): super().__init__(telemetry_messenger, POST_BREACH_PROCESS_LIST_COLLECTION) - def run(self): + def run(self, options: Dict): """ Collects process information from the host. Currently lists process name, ID, parent ID, command line 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 60990d67a..93b461c11 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 @@ -3,6 +3,7 @@ import random import shutil import string import subprocess +from typing import Dict from common.common_consts.post_breach_consts import POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER from infection_monkey.i_puppet.i_puppet import PostBreachData @@ -39,7 +40,7 @@ class CommunicateAsBackdoorUser(PBA): telemetry_messenger, name=POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER ) - def run(self): + def run(self, options: Dict): username = CommunicateAsBackdoorUser.get_random_new_user_name() try: password = get_random_password(14) diff --git a/monkey/infection_monkey/post_breach/actions/hide_files.py b/monkey/infection_monkey/post_breach/actions/hide_files.py index 457b9dafe..838fae222 100644 --- a/monkey/infection_monkey/post_breach/actions/hide_files.py +++ b/monkey/infection_monkey/post_breach/actions/hide_files.py @@ -1,3 +1,5 @@ +from typing import Dict + 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 @@ -21,7 +23,7 @@ class HiddenFiles(PBA): def __init__(self, telemetry_messenger: ITelemetryMessenger): super(HiddenFiles, self).__init__(telemetry_messenger, name=POST_BREACH_HIDDEN_FILES) - def run(self): + def run(self, options: Dict): # create hidden files and folders for function_to_get_commands in HIDDEN_FSO_CREATION_COMMANDS: linux_cmds, windows_cmds = function_to_get_commands() @@ -30,7 +32,7 @@ class HiddenFiles(PBA): linux_cmd=" ".join(linux_cmds), windows_cmd=windows_cmds, ) - super(HiddenFiles, self).run() + super(HiddenFiles, self).run(options) if is_windows_os(): # use winAPI result, status = get_winAPI_to_hide_files() 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 4d755567b..9b15de77f 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 @@ -1,4 +1,5 @@ import subprocess +from typing import Dict from common.common_consts.post_breach_consts import POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION from infection_monkey.i_puppet.i_puppet import PostBreachData @@ -19,7 +20,7 @@ class ModifyShellStartupFiles(PBA): def __init__(self, telemetry_messenger: ITelemetryMessenger): super().__init__(telemetry_messenger, name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION) - def run(self): + def run(self, options: Dict): results = [pba.run() for pba in self.modify_shell_startup_PBA_list()] if not results: results = [ diff --git a/monkey/infection_monkey/post_breach/actions/schedule_jobs.py b/monkey/infection_monkey/post_breach/actions/schedule_jobs.py index 8aeb0b42d..4ab023e35 100644 --- a/monkey/infection_monkey/post_breach/actions/schedule_jobs.py +++ b/monkey/infection_monkey/post_breach/actions/schedule_jobs.py @@ -1,3 +1,5 @@ +from typing import Dict + from common.common_consts.post_breach_consts import POST_BREACH_JOB_SCHEDULING from infection_monkey.post_breach.job_scheduling.job_scheduling import ( get_commands_to_schedule_jobs, @@ -22,7 +24,7 @@ class ScheduleJobs(PBA): windows_cmd=windows_cmds, ) - def run(self): - super(ScheduleJobs, self).run() + def run(self, options: Dict): + super(ScheduleJobs, self).run(options) remove_scheduled_jobs() return self.pba_data 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 d7323b54e..470e07bb1 100644 --- a/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py +++ b/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py @@ -1,5 +1,6 @@ import logging import subprocess +from typing import Dict from common.common_consts.post_breach_consts import POST_BREACH_SIGNED_SCRIPT_PROXY_EXEC from infection_monkey.post_breach.pba import PBA @@ -22,14 +23,14 @@ class SignedScriptProxyExecution(PBA): windows_cmd=" ".join(windows_cmds), ) - def run(self): + def run(self, options: Dict): original_comspec = "" try: if is_windows_os(): original_comspec = subprocess.check_output( # noqa: DUO116 "if defined COMSPEC echo %COMSPEC%", shell=True ).decode() - super().run() + super().run(options) return self.pba_data except Exception as e: logger.warning( diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 061fe1132..ec3f97134 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -37,7 +37,8 @@ class Puppet(IPuppet): return credential_collector.collect_credentials(options) def run_pba(self, name: str, options: Dict) -> Iterable[PostBreachData]: - return self._mock_puppet.run_pba(name, options) + pba = self._plugin_registry.get_plugin(name, PluginType.POST_BREACH_ACTION) + return pba.run(options) def ping(self, host: str, timeout: float = CONNECTION_TIMEOUT) -> PingScanData: return network_scanning.ping(host, timeout)