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 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
from infection_monkey.utils.environment import is_windows_os
@ -12,21 +12,31 @@ class ModifyShellStartupFiles(PBA):
and profile.ps1 in windows.
"""
def __init__(self):
super(ModifyShellStartupFiles, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION)
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()
if is_windows_os():
super(ModifyShellStartupFiles, self).__init__(name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION,
linux_cmd='', # windows so won't matter
windows_cmd=windows_cmds)
super(ModifyShellStartupFiles, self).run()
else:
for shell_startup_file in shell_startup_files_for_linux:
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,
windows_cmd=windows_cmds)
super(ModifyShellStartupFiles, self).run()
pbas = [ModifyShellStartupFile(linux_cmds='', windows_cmds=windows_cmds)]
for shell_startup_file in shell_startup_files_for_linux:
linux_cmds = ' '.join(cmds_for_linux).format(shell_startup_file)
pbas.append(ModifyShellStartupFile(linux_cmds=linux_cmds, windows_cmds=''))
return pbas
class ModifyShellStartupFile(PBA):
def __init__(self, linux_cmds, windows_cmds):
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
"""
exec_funct = self._execute_default
result = exec_funct()
if self.scripts_were_used_successfully(result):
T1064Telem(ScanStatus.USED, "Scripts were used to execute %s post breach action." % self.name).send()
PostBreachTelem(self, result).send()
if self.command:
exec_funct = self._execute_default
result = exec_funct()
if self.scripts_were_used_successfully(result):
T1064Telem(ScanStatus.USED, "Scripts were used to execute %s post breach action." % self.name).send()
PostBreachTelem(self, result).send()
def is_script(self):
"""