forked from p15670423/monkey
90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
import logging
|
|
import subprocess
|
|
|
|
from common.utils.attack_utils import ScanStatus
|
|
from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
|
|
from infection_monkey.utils import is_windows_os
|
|
from infection_monkey.config import WormConfiguration
|
|
from infection_monkey.telemetry.attack.t1064_telem import T1064Telem
|
|
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
|
|
__author__ = 'VakarisZ'
|
|
|
|
|
|
class PBA(object):
|
|
"""
|
|
Post breach action object. Can be extended to support more than command execution on target machine.
|
|
"""
|
|
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 = PBA.choose_command(linux_cmd, windows_cmd)
|
|
self.name = name
|
|
|
|
def get_pba(self):
|
|
"""
|
|
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.
|
|
"""
|
|
return self
|
|
|
|
@staticmethod
|
|
def should_run(class_name):
|
|
"""
|
|
Decides if post breach action is enabled in config
|
|
:return: True if it needs to be ran, false otherwise
|
|
"""
|
|
return class_name in WormConfiguration.post_breach_actions
|
|
|
|
def run(self):
|
|
"""
|
|
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()
|
|
|
|
def is_script(self):
|
|
"""
|
|
Determines if PBA is a script (PBA might be a single command)
|
|
:return: True if PBA is a script(series of OS commands)
|
|
"""
|
|
return isinstance(self.command, list) and len(self.command) > 1
|
|
|
|
def scripts_were_used_successfully(self, pba_execution_result):
|
|
"""
|
|
Determines if scripts were used to execute PBA and if they succeeded
|
|
:param pba_execution_result: result of execution function. e.g. self._execute_default
|
|
:return: True if scripts were used, False otherwise
|
|
"""
|
|
pba_execution_succeeded = pba_execution_result[1]
|
|
return pba_execution_succeeded and self.is_script()
|
|
|
|
def _execute_default(self):
|
|
"""
|
|
Default post breach command execution routine
|
|
:return: Tuple of command's output string and boolean, indicating if it succeeded
|
|
"""
|
|
try:
|
|
return subprocess.check_output(self.command, stderr=subprocess.STDOUT, shell=True), True
|
|
except subprocess.CalledProcessError as e:
|
|
# Return error output of the command
|
|
return e.output, False
|
|
|
|
@staticmethod
|
|
def choose_command(linux_cmd, windows_cmd):
|
|
"""
|
|
Helper method that chooses between linux and windows commands.
|
|
:param linux_cmd:
|
|
:param windows_cmd:
|
|
:return: Command for current os
|
|
"""
|
|
return windows_cmd if is_windows_os() else linux_cmd
|