Agent: Modify puppet to run PBAs instead of using the mock puppet

This commit is contained in:
Shreya Malviya 2022-03-30 13:37:47 +05:30
parent 394088e39d
commit 40b1ae0058
8 changed files with 21 additions and 11 deletions

View File

@ -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

View File

@ -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

View File

@ -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)

View File

@ -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()

View File

@ -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 = [

View File

@ -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

View File

@ -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(

View File

@ -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)