Code design changes

This commit is contained in:
Shreya 2020-06-17 11:49:03 +05:30
parent 60207e3a79
commit 0c60ad16ae
2 changed files with 32 additions and 21 deletions

View File

@ -1,6 +1,6 @@
from common.data.post_breach_consts import POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION from common.data.post_breach_consts import POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION
from infection_monkey.post_breach.pba import PBA from infection_monkey.post_breach.pba import PBA
from infection_monkey.utils.shell_startup_files_modification import\ from infection_monkey.post_breach.shell_startup_files.shell_startup_files_modification import\
get_commands_to_modify_shell_startup_files get_commands_to_modify_shell_startup_files
from infection_monkey.utils.environment import is_windows_os from infection_monkey.utils.environment import is_windows_os
@ -12,21 +12,31 @@ class ModifyShellStartupFiles(PBA):
and profile.ps1 in windows. and profile.ps1 in windows.
""" """
def __init__(self):
super(ModifyShellStartupFiles, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION)
def run(self): def run(self):
[pba.run() for pba in self.modify_shell_startup_PBA_list()]
def modify_shell_startup_PBA_list(self):
return ShellStartupPBAGenerator.get_modify_shell_startup_pbas()
class ShellStartupPBAGenerator():
def get_modify_shell_startup_pbas():
(cmds_for_linux, shell_startup_files_for_linux), windows_cmds = get_commands_to_modify_shell_startup_files() (cmds_for_linux, shell_startup_files_for_linux), windows_cmds = get_commands_to_modify_shell_startup_files()
if is_windows_os(): pbas = [ModifyShellStartupFile(linux_cmds='', windows_cmds=windows_cmds)]
super(ModifyShellStartupFiles, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION,
linux_cmd='', # windows so won't matter for shell_startup_file in shell_startup_files_for_linux:
windows_cmd=windows_cmds) linux_cmds = ' '.join(cmds_for_linux).format(shell_startup_file)
super(ModifyShellStartupFiles, self).run() pbas.append(ModifyShellStartupFile(linux_cmds=linux_cmds, windows_cmds=''))
else:
for shell_startup_file in shell_startup_files_for_linux: return pbas
linux_cmds = ' '.join(cmds_for_linux).format(shell_startup_file)
super(ModifyShellStartupFiles, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION,
linux_cmd=linux_cmds, class ModifyShellStartupFile(PBA):
windows_cmd=windows_cmds) def __init__(self, linux_cmds, windows_cmds):
super(ModifyShellStartupFiles, self).run() super(ModifyShellStartupFile, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION,
linux_cmd=linux_cmds,
windows_cmd=windows_cmds)
def run(self):
super(ModifyShellStartupFile, self).run()

View File

@ -57,11 +57,12 @@ class PBA(Plugin):
""" """
Runs post breach action command Runs post breach action command
""" """
exec_funct = self._execute_default if self.command:
result = exec_funct() exec_funct = self._execute_default
if self.scripts_were_used_successfully(result): result = exec_funct()
T1064Telem(ScanStatus.USED, "Scripts were used to execute %s post breach action." % self.name).send() if self.scripts_were_used_successfully(result):
PostBreachTelem(self, result).send() T1064Telem(ScanStatus.USED, "Scripts were used to execute %s post breach action." % self.name).send()
PostBreachTelem(self, result).send()
def is_script(self): def is_script(self):
""" """