integration of attackmate in external projects

README.md

@@ -67,3 +67,7 @@ a command injection. Keep that in mind!
 ## License
 
 [GPL-3.0](LICENSE)
+
+## Financial Support
+
+Funded by the European Union under GA no. 101121403 (NEWSROOM) and GA no. 101103385 (AInception). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Commission. Neither the European Union nor the granting authority can be held responsible for them. Further supported by the Horizon Europe project MIRANDA (101168144). Co-funded by the Austrian security research programme KIRAS of the Federal Ministry of Finance (BMF) in course of the projects ASOC (FO999905301) and Testcat (FO999911248).

docs/source/conf.py

@@ -13,7 +13,7 @@
 project = 'AttackMate'
 copyright = '2023, Wolfgang Hotwagner'
 author = 'Wolfgang Hotwagner'
-release = '0.2.1'
+release = '0.4.0'
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

docs/source/playbook/commands/debug.rst

@@ -3,7 +3,7 @@ debug
 =====
 
 This command prints out strings and variables and is for debugging
-purposes only.
+purposes/printing to the console only. This command does not modify the Builtin Variable ``RESULT_STDOUT``.
 
    .. code-block:: yaml
 

docs/source/playbook/commands/regex.rst

@@ -4,6 +4,7 @@ regex
 
 This command parses variables using regular expressions. For more information
 about regular expressions see `Python Regex  <https://docs.python.org/3/library/re.html>`_
+This command does not modify the Builtin Variable ``RESULT_STDOUT``.
 
 The following example parses the portnumber from the output of the last command and stores it in variable "UNREALPORT":
 

docs/source/playbook/commands/setvar.rst

@@ -4,6 +4,7 @@ setvar
 
 Set a variable. This could be used for string interpolation or for
 copying variables.
+This command does not modify the Builtin Variable ``RESULT_STDOUT``.
 
 .. code-block:: yaml
 

docs/source/playbook/vars.rst

@@ -41,7 +41,7 @@ Builtin Variables
 
 The following variables are set by the system:
 
-``RESULT_STDOUT`` is set after every command execution and stores the result output.
+``RESULT_STDOUT`` is set after every command execution (except for debug, regex and setvar commands) and stores the result output.
 
 ``RESULT_CODE`` is set after every command execution and stores the returncode.
 

src/attackmate/attackmate.py

@@ -8,19 +8,25 @@
 configuration.
 """
 
-from typing import Dict
+from typing import Dict, Optional
 import logging
+from attackmate.result import Result
 import attackmate.executors as executors
-from attackmate.schemas.config import Config
-from attackmate.schemas.playbook import Playbook, Commands
+from attackmate.schemas.config import CommandConfig, Config, MsfConfig, SliverConfig
+from attackmate.schemas.playbook import Playbook, Commands, Command
 from .variablestore import VariableStore
 from .processmanager import ProcessManager
 from attackmate.executors.baseexecutor import BaseExecutor
 from attackmate.executors.executor_factory import executor_factory
 
 
 class AttackMate:
-    def __init__(self, playbook: Playbook, config: Config) -> None:
+    def __init__(
+        self,
+        playbook: Optional[Playbook] = None,
+        config: Optional[Config] = None,
+        varstore: Optional[Dict] = None,
+    ) -> None:
         """Constructor for AttackMate
 
         This constructor initializes the logger('playbook'), the playbook,
@@ -33,20 +39,31 @@ def __init__(self, playbook: Playbook, config: Config) -> None:
         """
         self.logger = logging.getLogger('playbook')
         self.pm = ProcessManager()
-        self.pyconfig = config
-        self.playbook = playbook
-        self._initialize_variable_parser()
+
+        # Initialize playbook and config, with defaults if not provided
+        self.playbook = playbook if playbook else self._default_playbook()
+        self.pyconfig = config if config else self._default_config()
+
+        self._initialize_variable_parser(varstore)
         self.msfsessionstore = executors.MsfSessionStore(self.varstore)
         self.executor_config = self._get_executor_config()
         self.executors: Dict[str, BaseExecutor] = {}
 
-    def _initialize_variable_parser(self):
+    def _default_playbook(self) -> Playbook:
+        return Playbook(commands=[], vars={})
+
+    def _default_config(self) -> Config:
+        return Config(cmd_config=CommandConfig(), msf_config=MsfConfig(), sliver_config=SliverConfig())
+
+    def _initialize_variable_parser(self, varstore: Optional[Dict] = None):
         """Initializes the variable-parser
 
         The variablestore stores and replaces variables with values in certain strings
         """
         self.varstore = VariableStore()
-        self.varstore.from_dict(self.playbook.vars)
+        # if attackmate is imported and initialized in another project, vars can be passed as dict
+        # otherwise variable store is initializen with vars from playbook
+        self.varstore.from_dict(varstore if varstore else self.playbook.vars)
         self.varstore.replace_with_prefixed_env_vars()
 
     def _get_executor_config(self) -> dict:
@@ -76,6 +93,13 @@ def _run_commands(self, commands: Commands):
             if executor:
                 executor.run(command)
 
+    def run_command(self, command: Command) -> Result:
+        command_type = 'ssh' if command.type == 'sftp' else command.type
+        executor = self._get_executor(command_type)
+        if executor:
+            result = executor.run(command)
+        return result if result else Result(None, None)
+
     def main(self):
         """The main function
 

src/attackmate/command.py

@@ -0,0 +1,38 @@
+from typing import Dict, Tuple, Optional
+from attackmate.schemas.base import BaseCommand
+
+
+class CommandRegistry:
+    _type_registry: Dict[str, BaseCommand] = {}
+    _type_cmd_registry: Dict[Tuple[str, str], BaseCommand] = {}
+
+    @classmethod
+    def register(cls, type_: str, cmd: Optional[str] = None):
+        """register a command class by type or type + cmd."""
+
+        def decorator(command_class: BaseCommand):
+            if cmd:
+                cls._type_cmd_registry[(type_, cmd)] = command_class
+            else:
+                cls._type_registry[type_] = command_class
+            return command_class
+
+        return decorator
+
+    @classmethod
+    def get_command_class(cls, type_: str, cmd: Optional[str] = None):
+        """Retrieve the command class based on type or type + cmd."""
+        if cmd and (type_, cmd) in cls._type_cmd_registry:
+            return cls._type_cmd_registry[(type_, cmd)]
+        if type_ in cls._type_registry:
+            return cls._type_registry[type_]
+        raise ValueError(f"No command registered for type '{type_}' and cmd '{cmd}'")
+
+
+# Command interface
+class Command:
+    @staticmethod
+    def create(type: str, cmd: Optional[str] = None, **kwargs):
+        """return a command instance based on type and cmd."""
+        CommandClass = CommandRegistry.get_command_class(type, cmd)
+        return CommandClass(type=type, cmd=cmd, **kwargs)

src/attackmate/executors/baseexecutor.py

@@ -55,7 +55,7 @@ def __init__(
         self.output = logging.getLogger('output')
         self.substitute_cmd_vars = substitute_cmd_vars
 
-    def run(self, command: BaseCommand):
+    def run(self, command: BaseCommand) -> Result:
         """Execute the command
 
         This method is executed by AttackMate and
@@ -73,17 +73,18 @@ def run(self, command: BaseCommand):
 
         if command.only_if:
             if not Conditional.test(self.varstore.substitute(command.only_if, True)):
-                if hasattr(command, 'type'):
-                    self.logger.warning(f'Skipping {command.type}: {command.cmd}')
-                else:
-                    self.logger.warning(f'Skipping {command.cmd}')
-                return
+                self.logger.info(f'Skipping {getattr(command, "type", "")}({command.cmd})')
+                return Result(None, None)
         self.reset_run_count()
         self.logger.debug(f"Template-Command: '{command.cmd}'")
         if command.background:
-            self.exec_background(self.substitute_template_vars(command, self.substitute_cmd_vars))
+            # Background commands always return Result(None,None)
+            result = self.exec_background(self.substitute_template_vars(command, self.substitute_cmd_vars))
         else:
-            self.exec(self.substitute_template_vars(command, self.substitute_cmd_vars))
+            result = self.exec(self.substitute_template_vars(command, self.substitute_cmd_vars))
+
+        return result
+
 
     def log_command(self, command):
         """Log starting-status of the command"""
@@ -138,7 +139,7 @@ def save_output(self, command: BaseCommand, result: Result):
             except Exception as e:
                 self.logger.warning(f'Unable to write output to file {command.save}: {e}')
 
-    def exec(self, command: BaseCommand):
+    def exec(self, command: BaseCommand) -> Result:
         try:
             self.log_command(command)
             self.log_metadata(self.logger, command)
@@ -155,9 +156,11 @@ def exec(self, command: BaseCommand):
             self.error_if_or_not(command, result)
         self.loop_if(command, result)
         self.loop_if_not(command, result)
+        return result
 
-    def _loop_exec(self, command: BaseCommand):
-        self.exec(command)
+    def _loop_exec(self, command: BaseCommand) -> Result:
+        result = self.exec(command)
+        return result
 
     def _exec_cmd(self, command: Any) -> Result:
         return Result(None, None)

src/attackmate/executors/common/debugexecutor.py

@@ -19,12 +19,12 @@ def log_command(self, command: DebugCommand):
             self.logger.warning(self.varstore.variables)
 
     def _exec_cmd(self, command: DebugCommand) -> Result:
-        self.setoutuptvars = False
+        self.setoutputvars = False
         ret = 0
         if command.wait_for_key:
             self.logger.warning("Type enter to continue")
             input()
         if command.exit:
             ret = 1
 
-        return Result('', ret)
+        return Result(f"Debug: '{command.cmd}'", ret)

src/attackmate/executors/common/regexexecutor.py

@@ -60,7 +60,7 @@ def forge_and_register_variables(self, output: dict, data):
         return
 
     def _exec_cmd(self, command: RegExCommand) -> Result:
-        self.setoutuptvars = False
+        self.setoutputvars = False
         if command.mode == 'findall':
             m = re.findall(command.cmd, self.varstore.get_str(command.input))
             self.forge_and_register_variables(command.output, m)

src/attackmate/executors/common/setvarexecutor.py

@@ -33,7 +33,7 @@ def log_command(self, command: SetVarCommand):
         self.logger.warning(f"Setting Variable: '{command.variable}'")
 
     def _exec_cmd(self, command: SetVarCommand) -> Result:
-        self.setoutuptvars = False
+        self.setoutputvars = False
         content = command.cmd
         if command.encoder:
             try:

src/attackmate/executors/features/background.py

@@ -28,23 +28,20 @@ def __init__(self, pm: ProcessManager):
     def _create_queue(self) -> Optional[Queue]:
         return None
 
-    def exec_background(self, command: BaseCommand):
-        if hasattr(command, 'type'):
-            self.logger.info(f'Run in background: {command.type}({command.cmd})')
-        else:
-            self.logger.info(f'Run in background: {command.cmd}')
+    def exec_background(self, command: BaseCommand) -> Result:
+        self.logger.info(f'Run in background: {getattr(command, "type", "")}({command.cmd})')
 
         queue = self._create_queue()
 
         if queue:
-            p = self.pm.ctx.Process(target=self._exec_bg_cmd,
-                                    args=(command, queue))
+            p = self.pm.ctx.Process(target=self._exec_bg_cmd, args=(command, queue))
         else:
-            p = self.pm.ctx.Process(target=self._exec_bg_cmd,
-                                    args=(command,))
+            p = self.pm.ctx.Process(target=self._exec_bg_cmd, args=(command,))
         p.start()
         p.join(5)
         self.pm.add_process(p, command.kill_on_exit)
+        # background commands always return None Result
+        return Result(None, None)
 
     def _exec_bg_cmd(self, command: Any, queue: Optional[Queue] = None):
         self.is_child_proc = True

src/attackmate/executors/features/cmdvars.py

@@ -9,10 +9,10 @@
 class CmdVars:
     def __init__(self, variablestore: VariableStore):
         self.varstore = variablestore
-        self.setoutuptvars = True
+        self.setoutputvars = True
 
     def set_result_vars(self, result: Result):
-        if self.setoutuptvars:
+        if self.setoutputvars:
             self.varstore.set_variable('RESULT_STDOUT', result.stdout)
             self.varstore.set_variable('RESULT_RETURNCODE', str(result.returncode))
 

src/attackmate/executors/features/looper.py

@@ -14,42 +14,36 @@ def __init__(self, cmdconfig: CommandConfig):
         self.reset_run_count()
 
     def _loop_exec(self, command: BaseCommand):
-        if hasattr(command, 'type'):
-            err_str = f'loop_exec for command {command.type} is not implemented!'
-        else:
-            err_str = 'loop_exec for command is not implemented!'
+        err_str = f'loop_exec for command {getattr(command, "type", "")} is not implemented!'
         self.logger.error(err_str)
         exit(1)
 
     def reset_run_count(self):
         self.run_count = 1
 
     def loop_if(self, command: BaseCommand, result: Result):
-        if command.loop_if is not None:
-            m = re.search(command.loop_if, result.stdout, re.MULTILINE)
-            if m is not None:
-                self.logger.warning(f'Re-run command because loop_if matches: {m.group(0)}')
-                if self.run_count < CmdVars.variable_to_int('loop_count', command.loop_count):
-                    self.run_count = self.run_count + 1
-                    time.sleep(self.cmdconfig.loop_sleep)
-                    self._loop_exec(command)
-                else:
-                    self.logger.error('Exiting because loop_count exceeded')
-                    exit(1)
+        if command.loop_if and (match := re.search(command.loop_if, result.stdout, re.MULTILINE)):
+            self.logger.warning(f'Re-run command because loop_if matches: {match.group(0)}')
+            if self.run_count < CmdVars.variable_to_int('loop_count', command.loop_count):
+                self.run_count += 1
+                time.sleep(self.cmdconfig.loop_sleep)
+                self._loop_exec(command)
             else:
-                self.logger.debug('loop_if does not match')
+                self.logger.error('Exiting because loop_count exceeded')
+                exit(1)
+        else:
+            self.logger.debug('loop_if does not match')
 
     def loop_if_not(self, command: BaseCommand, result: Result):
-        if command.loop_if_not is not None:
-            m = re.search(command.loop_if_not, result.stdout, re.MULTILINE)
-            if m is None:
-                self.logger.warning('Re-run command because loop_if_not does not match')
-                if self.run_count < CmdVars.variable_to_int('loop_count', command.loop_count):
-                    self.run_count = self.run_count + 1
-                    time.sleep(self.cmdconfig.loop_sleep)
-                    self._loop_exec(command)
-                else:
-                    self.logger.error('Exitting because loop_count exceeded')
-                    exit(1)
+        if command.loop_if_not and re.search(command.loop_if_not, result.stdout, re.MULTILINE) is None:
+            self.logger.warning('Re-run command because loop_if_not does not match')
+
+            if self.run_count < CmdVars.variable_to_int('loop_count', command.loop_count):
+                self.run_count += 1
+                time.sleep(self.cmdconfig.loop_sleep)
+                self._loop_exec(command)
             else:
-                self.logger.debug('loop_if_not does not match')
+                self.logger.error('Exiting because loop_count exceeded')
+                exit(1)
+        else:
+            self.logger.debug('loop_if_not does not match')

src/attackmate/metadata.py

@@ -8,7 +8,7 @@
 __license__ = 'GPLv3'
 __maintainer__ = 'Wolfgang Hotwagner, Max Landauer, Markus Wurzenberger, Florian Skopik'
 __status__ = 'Production'
-__version__ = '0.2.1'
+__version__ = '0.4.0'
 __version_string__ = (
         f'(Austrian Institute of Technology)\t'
         f'{__website__}\tVersion: {__version__}')

src/attackmate/result.py

@@ -24,3 +24,7 @@ def __init__(self, stdout, returncode):
         """
         self.stdout = stdout
         self.returncode = returncode
+
+    def __repr__(self):
+            return f"Result(stdout={repr(self.stdout)}, returncode={self.returncode})"
+

src/attackmate/schemas/debug.py

@@ -1,7 +1,9 @@
 from typing import Literal
 from .base import BaseCommand
+from attackmate.command import CommandRegistry
 
 
+@CommandRegistry.register('debug')
 class DebugCommand(BaseCommand):
     type: Literal['debug']
     varstore: bool = False