diff --git a/docs/source/developing/baseexecutor.rst b/docs/source/developing/baseexecutor.rst index 0d5f7e16..e156905c 100644 --- a/docs/source/developing/baseexecutor.rst +++ b/docs/source/developing/baseexecutor.rst @@ -31,6 +31,7 @@ Implementing a Custom Executor ------------------------------- To create a custom executor, inherit from ``BaseExecutor`` and implement the ``_exec_cmd()`` method. Other methods can be overriden as needed. +New executors must be registered using the @executor_factory.register_executor('') decorator. Example: @@ -99,4 +100,20 @@ executor __init__.py -------------------- .. note:: - Add the new executor to the ``__all__`` list in the ``__init__.py`` file of the ``attackmate.executors`` module. + Add the new executor to the ``__all__`` list in the ``__init__.py`` file of the ``attackmate.executors`` module so it can be imported elsewhere. + +.. code-block:: python + + # src/attackmate/executors/__init__.py + # other imports + from .shell.shellexecutor import ShellExecutor + from .metasploit.msfsessionexecutor import CustomExecutor # new executor + # other imports + + __all__ = [ + 'RemoteExecutor', + 'BrowserExecutor', + 'ShellExecutor', + 'CustomExecutor', # new executor + # other executors + ] diff --git a/docs/source/developing/command.rst b/docs/source/developing/command.rst index 18d75e59..9896d4a3 100644 --- a/docs/source/developing/command.rst +++ b/docs/source/developing/command.rst @@ -33,7 +33,7 @@ make it usable in external python scripts. **Example: a simple command with no sub-behaviors** -:: +.. code-block:: python from typing import Literal from .base import BaseCommand @@ -50,7 +50,7 @@ make it usable in external python scripts. **Example: a command with sub-behaviors expressed via** ``cmd`` -:: +.. code-block:: python from typing import Literal from .base import BaseCommand @@ -68,7 +68,7 @@ make it usable in external python scripts. The new command should be handled by an executor in `src/attackmate/executors` that extends ``BaseExecutor`` and implements the ``_exec_cmd()`` method. For example: -:: +.. code-block:: python from attackmate.executors.base_executor import BaseExecutor from attackmate.result import Result @@ -92,7 +92,7 @@ When a command is executed, the ``create_executor`` method retrieves the corresp Accordingly, executors must be registered using the ``@executor_factory.register_executor('')`` decorator. -:: +.. code-block:: python @executor_factory.register_executor('debug') class DebugExecutor(BaseExecutor): @@ -103,7 +103,7 @@ If the new executor class requires additional initialization arguments, these mu All configurations are always passed to the ``ExecutorFactory``. The factory filters the provided configurations based on the class constructor signature, ensuring that only the required parameters are used. -:: +.. code-block:: python def _get_executor_config(self) -> dict: config = { @@ -118,13 +118,35 @@ The factory filters the provided configurations based on the class constructor s } return config +4. Add the Executor to the __init__ file of the attackmate.executors module +=========================================================================== -4. Modify the Loop Command to Include the New Command +Add the new executor to the __all__ list in the __init__.py file of the attackmate.executors module so it can be imported elsewhere. + +.. code-block:: python + + # src/attackmate/executors/__init__.py + # other imports + from .shell.shellexecutor import ShellExecutor + from .metasploit.msfsessionexecutor import CustomExecutor # new executor + # other imports + + __all__ = [ + 'RemoteExecutor', + 'BrowserExecutor', + 'ShellExecutor', + 'CustomExecutor', # new executor + # other executors + ] + + + +5. Modify the Loop Command to Include the New Command ===================================================== in `/src/attackmate/schemas/loop.py` update the ``LoopCommand`` schema to include the new command. -:: +.. code-block:: python Command = Union[ ShellCommand, @@ -133,11 +155,12 @@ in `/src/attackmate/schemas/loop.py` update the ``LoopCommand`` schema to includ ] -4. Modify the RemotelyExecutableCommand Union to Include the New Command +6. Modify the RemotelyExecutableCommand Union to Include the New Command ======================================================================== in `src/attackmate/schemas/command_subtypes.py`, update the ``RemotelyExecutableCommand`` type alias to include the new command -:: + +.. code-block:: python RemotelyExecutableCommand: TypeAlias = Annotated[ Union[ @@ -176,7 +199,7 @@ in `src/attackmate/schemas/command_subtypes.py`, update the ``RemotelyExecutable Once these steps are completed, the new command will be fully integrated into AttackMate and available for execution. -6. Add Documentation +7. Add Documentation ===================== Finally, update the documentation in `docs/source/playbook/commands` to include the new command. diff --git a/docs/source/developing/integration.rst b/docs/source/developing/integration.rst index 4bb22d52..82d2c40e 100644 --- a/docs/source/developing/integration.rst +++ b/docs/source/developing/integration.rst @@ -67,11 +67,11 @@ Understanding the Result Object =============================== When executing a command with AttackMate, the result is returned as an instance of the ``Result`` class. This object contains the standard output (`stdout`) and the return code (`returncode`) of the executed command. -Commands that run in the Background return Result('Command started in background', 0) +Commands that run in the :ref:`background` return Result('Command started in background', 0) .. note:: Regular Commands return a ``Result`` object. - Commands that run in background mode return ``Result('Command started in background', 0)``. + Commands that run in :ref:`background` mode return ``Result('Command started in background', 0)``. Attributes ---------- diff --git a/docs/source/index.rst b/docs/source/index.rst index 9267fa7b..00b48db0 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -4,7 +4,7 @@ contain the root `toctree` directive. Welcome to AttackMate's documentation! -================================== +====================================== .. toctree:: :maxdepth: 1 @@ -32,6 +32,7 @@ Welcome to AttackMate's documentation! playbook/commands/index playbook/session/index playbook/examples + playbook/troubleshooting .. toctree:: :maxdepth: 1 diff --git a/docs/source/playbook/commands/browser.rst b/docs/source/playbook/commands/browser.rst index 3733e195..76151c02 100644 --- a/docs/source/playbook/commands/browser.rst +++ b/docs/source/playbook/commands/browser.rst @@ -1,3 +1,5 @@ +.. _browser: + ======= browser ======= @@ -11,6 +13,9 @@ Execute commands using a Playwright-managed Chromium browser. This executor can etc.) across multiple commands, use the :confval:`creates_session` to open a named session and :confval:`session` to reuse it. + Background mode is not supported for this commands. + + .. code-block:: yaml vars: diff --git a/docs/source/playbook/commands/include.rst b/docs/source/playbook/commands/include.rst index fd2cb66b..497f58cc 100644 --- a/docs/source/playbook/commands/include.rst +++ b/docs/source/playbook/commands/include.rst @@ -1,9 +1,15 @@ +.. _include: + ======= include ======= Include and run commands from another yaml-file. +.. note:: + + Background mode is not supported for this commands. + .. code-block:: yaml # main.yml: diff --git a/docs/source/playbook/commands/index.rst b/docs/source/playbook/commands/index.rst index f0909163..9454fd3f 100644 --- a/docs/source/playbook/commands/index.rst +++ b/docs/source/playbook/commands/index.rst @@ -28,6 +28,24 @@ Every command, regardless of its type supports the following general options: cmd: nmap localhost save: /tmp/nmap_localhost.txt +.. confval:: metadata + + An optional dictionary of key-value pairs that are logged alongside the command + but have no effect on execution. + + :type: Dict + :default: None + + .. code-block:: yaml + + commands: + - type: debug + cmd: Come on, Cat + metadata: + version: 1 + author: Ellen Ripley + + .. confval:: exit_on_error Attackmate will exit with an error if the command returns a non-zero exit code. @@ -107,6 +125,8 @@ Every command, regardless of its type supports the following general options: :type: int :default: ``3`` +.. _conditionals: + .. confval:: only_if Execute this command only if the condition evaluates to ``True``. Supported operators: @@ -153,6 +173,8 @@ Every command, regardless of its type supports the following general options: .. warning:: When comparing strings with integers, standard Python conventions apply: + see `Python reference — Comparisons + `_ **Equality / Inequality** (``==``, ``!=``): A string and an integer are never equal, so ``"1" == 1`` is ``False`` @@ -161,7 +183,8 @@ Every command, regardless of its type supports the following general options: **Identity** (``is``, ``is not``): Compares object identity, not value. ``"1" is 1`` is always ``False`` because a string and an integer are distinct objects, regardless of - their apparent values. + their apparent values. see `Python reference — Value vs. identity + `_ **Ordering** (``<``, ``<=``, ``>``, ``>=``): Comparing a string with an integer raises a ``TypeError`` in Python 3 @@ -178,12 +201,18 @@ Every command, regardless of its type supports the following general options: while the literal is parsed as a ``bool`` by ``ast``. Use string literals for boolean-like flags stored in the ``VariableStore``: ``$flag == "True"``. + see `Python built-in — bool (subclass of int) + `_ Importantly, before a condition is evaluated, all ``$variable`` references are resolved by the ``VariableStore``. **The store holds every value as a plain Python** ``str``, even values that were originally integers are coerced to ``str`` on ingress. + +.. _background: + + .. confval:: background Execute the command as a background subprocess. When enabled, output is not printed @@ -194,23 +223,20 @@ Every command, regardless of its type supports the following general options: .. note:: - The command in background-mode will change the :ref:`builtin variables ` + The command in background mode will change the :ref:`builtin variables ` ``RESULT_STDOUT`` to "Command started in Background" and ``RESULT_CODE`` to 0. Background mode is not supported for - * MsfModuleCommand - * IncludeCommand - * VncCommand - * BrowserCommand + * :ref:`MsfModuleCommand ` + * :ref:`IncludeCommand ` + * :ref:`VncCommand ` + * :ref:`BrowserCommand ` Background mode together with a session is not supported for the following commands: - * SSHCommand - * SFTPCommand - - - + * :ref:`SSHCommand ` + * :ref:`SFTPCommand ` .. confval:: kill_on_exit @@ -221,23 +247,6 @@ Every command, regardless of its type supports the following general options: :type: bool :default: ``True`` -.. confval:: metadata - - An optional dictionary of key-value pairs that are logged alongside the command - but have no effect on execution. - - :type: Dict - :default: None - - .. code-block:: yaml - - commands: - - type: debug - cmd: Come on, Cat - metadata: - version: 1 - author: Ellen Ripley - The next pages will describe each command type in detail. diff --git a/docs/source/playbook/commands/msf-module.rst b/docs/source/playbook/commands/msf-module.rst index e10f9e12..2fd2f474 100644 --- a/docs/source/playbook/commands/msf-module.rst +++ b/docs/source/playbook/commands/msf-module.rst @@ -10,6 +10,9 @@ Execute Metasploit modules via the Metasploit RPC API. To configure the connection to ``msfrpcd``, see :ref:`msf_config`. + Background mode is not supported for this commands. + + Some modules (like auxiliary scanners) produce direct output: .. code-block:: yaml diff --git a/docs/source/playbook/commands/sftp.rst b/docs/source/playbook/commands/sftp.rst index 0d5839e1..d991e0c5 100644 --- a/docs/source/playbook/commands/sftp.rst +++ b/docs/source/playbook/commands/sftp.rst @@ -1,3 +1,5 @@ +.. _sftp: + ==== sftp ==== @@ -8,8 +10,11 @@ created by either command can be reused by the other. .. note:: - This command caches all settings so - that they only need to be defined once. + This command caches all settings so that they only need to be defined once. + + Background mode with a session is not supported for this commands. + + .. code-block:: yaml vars: @@ -45,8 +50,8 @@ File Transfer The SFTP operation to perform. - * ``put`` — upload a file from the local machine to the remote host - * ``get`` — download a file from the remote host to the local machine + * ``put`` - upload a file from the local machine to the remote host + * ``get`` - download a file from the remote host to the local machine :type: str :required: True diff --git a/docs/source/playbook/commands/ssh.rst b/docs/source/playbook/commands/ssh.rst index a36e7a4c..57b24c94 100644 --- a/docs/source/playbook/commands/ssh.rst +++ b/docs/source/playbook/commands/ssh.rst @@ -11,6 +11,8 @@ Execute commands on a remote host via SSH. This command caches all settings so that they only need to be defined once. + Background mode with a session is not supported for this commands. + .. code-block:: yaml vars: diff --git a/docs/source/playbook/commands/vnc.rst b/docs/source/playbook/commands/vnc.rst index 7dd3b80a..ce6d44fe 100644 --- a/docs/source/playbook/commands/vnc.rst +++ b/docs/source/playbook/commands/vnc.rst @@ -1,3 +1,5 @@ +.. _vnc: + === vnc === @@ -11,6 +13,10 @@ cached after the first command, so subsequent commands only need to specify what VNC sessions must be explicitly closed with ``cmd: close``, otherwise AttackMate will hang on exit. +.. note:: + + Background mode is not supported for this commands. + .. code-block:: yaml vars: diff --git a/docs/source/playbook/troubleshooting.rst b/docs/source/playbook/troubleshooting.rst new file mode 100644 index 00000000..c940df89 --- /dev/null +++ b/docs/source/playbook/troubleshooting.rst @@ -0,0 +1,158 @@ +.. _troubleshooting: + +=============== +Troubleshooting +=============== + +This page lists common error messages and their solutions. + +.. _error-validation: + +Playbook Validation Errors +========================== + +These errors occur when AttackMate cannot parse the playbook file. +They are reported as ``ValidationError`` and originate from Pydantic's +model validation. + +---- + +.. _error-invalid-command-type: + +Invalid Command Type +-------------------- + +**Symptom** + +.. code-block:: text + + ERROR | A Validation error occured when parsing playbook file playbooks/example.yml + ERROR | Traceback (most recent call last): + ... + pydantic_core._pydantic_core.ValidationError: 1 validation error for Playbook + commands.0 + Input tag 'shel' found using 'type' does not match any of the expected tags: + 'sliver-session', 'sliver', 'browser', 'shell', 'msf-module', 'msf-session', + 'msf-payload', 'sleep', 'ssh', 'father', 'sftp', 'debug', 'setvar', 'regex', + 'mktemp', 'include', 'loop', 'webserv', 'http-client', 'json', 'vnc', + 'bettercap', 'remote' + [type=union_tag_invalid, input_value={'type': 'shel', 'cmd': 'whoami'}, input_type=dict] + +**Cause** + +The ``type`` field of a command does not match any registered command type. +In the example above, ``shel`` is a typo for ``shell``. + +**Solution** + +* Check the spelling of the ``type`` field in your playbook. Valid types are + listed in the error message itself. +* If this is a new command type you are trying to implement, verify that it is correctly + registered in the :class:`CommandRegistry` and included in the executor + factory. see :ref:`Adding a new Command `. + +.. seealso:: + see :ref:`commands` for a full list of available command types. + +---- + +.. _error-missing-field: + +Missing Required Field +---------------------- + +**Symptom** + +.. code-block:: text + + ERROR | A Validation error occured when parsing playbook file playbooks/example.yml + ERROR | Missing field in shell command: cmd - Field required + ERROR | Traceback (most recent call last): + ... + pydantic_core._pydantic_core.ValidationError: 1 validation error for Playbook + commands.0.shell.cmd + Field required [type=missing, input_value={'type': 'shell'}, input_type=dict] + +**Cause** + +A required field is missing from a command definition. In the example above, +the ``cmd`` field is missing from a ``shell`` command. + +**Solution** + +* Check the command definition in your playbook and ensure all required fields + are present. +* Refer to the documentation for the specific command type to see which fields + are required. + +---- + +Sliver Command Execution Errors +=============================== + +These errors occor when there is a faulty sliver configuration, sliver is not installed etc. + +.. _error-sliver-client-not-defined: + +Sliver Client Not Defined +------------------------- + +**Symptom** + +.. code-block:: text + + INFO | Executing Sliver-command: 'start_https_listener' + ERROR | Sliver execution failed: SliverClient is not defined + ERROR | SliverClient is not defined + +**Cause** + +A Sliver command was executed but no ``SliverClient`` was defined. + +**Solution** + +* Verify that the ``sliver_config`` in your configuration points to a valid + Sliver client configuration file and that sliver is installed. + +.. seealso:: + :ref:`sliver` and :ref:`sliver_config` for setup instructions and required configuration fields. + +---- + +.. _error-sliver-connection-refused: + +Sliver Server Unreachable +------------------------- + +**Symptom** + +.. code-block:: text + + INFO | Executing Sliver-command: 'start_https_listener' + ERROR | An error occurred: + +**Cause** + +AttackMate could not establish a gRPC connection to the Sliver C2 server. +The server is either not running, not reachable at the configured address +and port, or is blocked by a firewall. + +**Solution** + +* Verify that the Sliver C2 server is running, the command to start it is ``sliver-server``, + per default it listens on port ``31337`` +* Check that the ``config_file`` in ``sliver_config`` points to the correct + client configuration file and that the host and port match the running + server (default: ``127.0.0.1:31337``). +* Ensure no firewall or network policy is blocking the connection. + +.. seealso:: + :ref:`sliver` and :ref:`sliver_config` for setup instructions and required configuration fields. diff --git a/docs/source/playbook/vars.rst b/docs/source/playbook/vars.rst index 335c133c..5360cb17 100644 --- a/docs/source/playbook/vars.rst +++ b/docs/source/playbook/vars.rst @@ -9,30 +9,121 @@ placeholders in command settings. Variable names do not require a ``$`` prefix w defined in the ``vars`` section, but **MUST** be prefixed with ``$`` when referenced in the ``commands`` section. If an environment variable with the prefix ``ATTACKMATE_`` exists with the same name, -it will override the playbook variable. For example, the playbook variabel ``$FOO`` will -be overwritten by the environment variabel ``$ATTACKMATE_FOO``. +it will override the playbook variable. For example, the playbook variable ``$FOO`` will +be overwritten by the environment variable ``$ATTACKMATE_FOO``. .. code-block:: yaml - vars: - # the $-sign is optional here: - $SERVER_ADDRESS: 192.42.0.254 - $NMAP: /usr/bin/nmap + vars: + # the $-sign is optional here: + $SERVER_ADDRESS: 192.42.0.254 + $NMAP: /usr/bin/nmap - commands: - - type: shell - # the $-sign is required when referencing a variable: - cmd: $NMAP $SERVER_ADDRESS + commands: + - type: shell + # the $-sign is required when referencing a variable: + cmd: $NMAP $SERVER_ADDRESS .. note:: - Variable substitution uses Python's `string.Template `_ syntax. - .. note:: + Variables in ``cmd`` settings of a ``loop`` command will be substituted on every + iteration of the loop, see the :ref:`loop` command for details. + + +.. _variable-types: + +Variable Types and Storage +========================== + +All variables are stored internally as plain Python ``str`` values, regardless of +how they were originally defined. Integer values are coerced to strings on ingress — +``set_variable`` explicitly converts ``int`` to ``str`` before storing. This means +that even if a variable is set programmatically to the integer ``42``, it will be +stored and retrieved as the string ``"42"``. + +The store holds two separate namespaces: + +- **Scalar variables** — simple ``str`` key-value pairs, accessed as ``$varname``. +- **List variables** — ordered sequences of ``str`` values, accessed by index as + ``$varname[0]``, ``$varname[1]``, etc. + +A name cannot be both a scalar and a list simultaneously. Assigning a list value to +an existing scalar name (or vice versa) silently replaces the previous entry. + + +.. _variable-comparison: + +Comparing Variables +=================== + +When a variable is used in a conditional expression (see :ref:`conditionals`), its +``$`` reference is substituted *before* the expression is evaluated. Because the +store always holds strings, the resolved value will be a ``str`` — even if the +original value looked like a number or a boolean. + +This has important consequences depending on the operator used: + +**Strings and integers** (``==``, ``!=``, ``<``, ``<=``, ``>``, ``>=``, ``is``, ``is not``) + The resolved variable is always a ``str``. If the right-hand side of the condition + is written as a bare integer literal (e.g. ``$exit_code == 0``), Python will compare + a ``str`` against an ``int``. Always use string literals on the right-hand side: + + .. code-block:: yaml + + # Wrong — "0" == 0 is False in Python: + condition: $exit_code == 0 + + # Correct: + condition: $exit_code == "0" + + Ordering operators (``<``, ``<=``, ``>``, ``>=``) between a ``str`` and an ``int`` + raise a ``TypeError``. If both sides are string literals the comparison uses + **lexicographic** ordering (``"9" > "10"`` is ``True``), so take care when + comparing numeric strings. + +**Strings and booleans** (``==``, ``!=``) + In Python, ``bool`` is a subclass of ``int``, so ``True`` equals ``1`` and + ``False`` equals ``0``. However, a variable set to ``"True"`` or ``"False"`` + is a ``str``, and ``"True" == True`` is ``False`` because the types differ. + Use string literals for boolean-like flags: + + .. code-block:: yaml + + # Wrong — "True" == True is False: + condition: $flag == True + + # Correct: + condition: $flag == "True" + +**Identity** (``is``, ``is not``) + These operators test object identity, not value equality. Because every resolved + variable is a freshly substituted ``str``, ``$var is "hello"`` is unreliable even + when the values appear equal. Use ``==`` and ``!=`` for value comparisons; reserve + ``is`` / ``is not`` for checks between two variables where identity is + intentionally meaningful. + +**Regex** (``=~``, ``!~``) + Regex conditions bypass ``ast`` parsing entirely and operate directly on the + substituted string, so type-mismatch issues do not apply. This is the most + robust operator for testing variable values that may contain numbers, booleans, + or mixed content: + + .. code-block:: yaml + + condition: $exit_code =~ ^0$ + +.. seealso:: + `Python reference — Comparisons + `_ + + `Python built-in — bool (subclass of int) + `_ - Variables in ``cmd`` settings of a ``loop`` command will be substituted on every iteration of the loop, see the :ref:`loop` command for details. + `Python ast — Constant node + `_ .. _builtin-variables: @@ -42,14 +133,22 @@ Builtin Variables The following variables are set automatically by AttackMate during execution: -``RESULT_STDOUT`` Stores the standard output of the most recently executed command. Not set by ``debug``, ``regex``, or ``setvar`` commands. +``RESULT_STDOUT`` + Stores the standard output of the most recently executed command. + Not set by ``debug``, ``regex``, or ``setvar`` commands. -``RESULT_CODE`` Stores the return code of the most recently executed command. +``RESULT_CODE`` + Stores the return code of the most recently executed command. -``LAST_MSF_SESSION`` Set whenever a new Metasploit session is created. Contains the session number. +``LAST_MSF_SESSION`` + Set whenever a new Metasploit session is created. Contains the session number. -``LAST_SLIVER_IMPLANT`` Set whenever a new Sliver implant is generated. Contains the path to the implant file. +``LAST_SLIVER_IMPLANT`` + Set whenever a new Sliver implant is generated. Contains the path to the implant file. -``LAST_FATHER_PATH`` Set whenever a Father rootkit is generated. Contains the path to the rootkit. +``LAST_FATHER_PATH`` + Set whenever a Father rootkit is generated. Contains the path to the rootkit. -``REGEX_MATCHES_LIST`` Set every time a regex command yields matches. Contains a list of all matches. If ``sub`` or ``split`` finds no match, the original input string is returned. +``REGEX_MATCHES_LIST`` + Set every time a regex command yields matches. Contains a list of all matches. + If ``sub`` or ``split`` finds no match, the original input string is returned. diff --git a/src/attackmate/attackmate.py b/src/attackmate/attackmate.py index 082653ac..292a94a8 100644 --- a/src/attackmate/attackmate.py +++ b/src/attackmate/attackmate.py @@ -1,11 +1,17 @@ -"""AttackMate reads a playbook and executes the attack - -A playbook stored in a dictionary with a list of attacks. Attacks -are executed by "Executors". There are many different -Executors like: ShellExecutor, SleepExecutor or MsfModuleExecutor -This class creates instances of all possible Executors, iterates -over all attacks and runs the specific Executor with the given -configuration. +"""AttackMate reads a playbook and executes the attack chain. + +A playbook is a structured sequence of commands, each dispatched to the +appropriate executor based on its type. Executors are instantiated lazily +on first use and cached for reuse. Available executors include +:class:`ShellExecutor`, :class:`SleepExecutor`, :class:`MsfModuleExecutor`, +and others registered via the executor factory. + +This module exposes :class:`AttackMate`, which can be used standalone via +its :meth:`~AttackMate.main` entry point or embedded in a Python script +using :meth:`~AttackMate.run_command` for single-command execution. + +.. seealso:: + :ref:`integration` for documentation on scripted usage. """ import time @@ -23,13 +29,26 @@ class AttackMate: - """ - Reads a playbook and executes the attack chain. + """Reads a playbook and executes the attack chain. + + Creates instances of all registered executors, iterates over the commands + in the playbook, and dispatches each to the appropriate executor. - :param playbook: The playbook to execute. - :param config: AttackMate configuration. - :param varstore: Initial variable store. - :param is_api_instance: Whether this instance is used as an API. + Can be used standalone via :meth:`main`, or embedded in a Python script + using :meth:`run_command` for single-command execution. + + :param playbook: The playbook to execute. Defaults to an empty playbook. + :param config: AttackMate configuration. Defaults to a default :class:`Config`. + :param varstore: Initial variables passed as a plain dictionary. If omitted, + variables are read from the playbook's ``vars`` section. + :param is_api_instance: Set to ``True`` when AttackMate is used as an embedded + API rather than a standalone CLI process. + + Example:: + + attackmate = AttackMate(config=config, varstore={"HOST": "10.0.0.1"}) + command = Command.create(type="shell", cmd="whoami") + result = await attackmate.run_command(command) """ def __init__( @@ -81,7 +100,7 @@ def _initialize_variable_parser(self, varstore: Optional[Dict] = None): """ self.varstore = VariableStore() # if attackmate is imported and initialized in another project, vars can be passed as dict - # otherwise variable store is initializen with vars from playbook + # otherwise variable store is initialized with vars from playbook self.varstore.from_dict(varstore if varstore else self.playbook.vars) self.varstore.replace_with_prefixed_env_vars() @@ -100,6 +119,15 @@ def _get_executor_config(self) -> dict: return config def _get_executor(self, command_type: str) -> BaseExecutor: + """Return the executor instance for the given command type. + + Executors are instantiated lazily on first use and cached for subsequent + calls. If no executor for ``command_type`` exists yet, one is created via + :func:`executor_factory.create_executor` and stored for reuse. + + :param command_type: The command type string (e.g. ``"shell"``, ``"ssh"``). + :returns: The :class:`BaseExecutor` instance for that command type. + """ if command_type not in self.executors: self.executors[command_type] = executor_factory.create_executor( command_type, **self.executor_config @@ -108,6 +136,21 @@ def _get_executor(self, command_type: str) -> BaseExecutor: return self.executors[command_type] async def _run_commands(self, commands: Commands): + """Execute a sequence of commands in order. + + Iterates over ``commands``, resolves the appropriate executor for each, + and runs them sequentially. A configurable delay is applied before each + command, except for ``sleep``, ``debug``, and ``setvar`` commands which + are exempt from the delay. + + ``sftp`` commands are dispatched to the ``ssh`` executor. + + :param commands: A sequence of command instances to execute. + + .. note:: + The delay between commands is controlled by ``cmd_config.command_delay`` + in the :class:`Config`. Defaults to ``0`` if not set. + """ delay = self.pyconfig.cmd_config.command_delay or 0 self.logger.info(f'Delay before commands: {delay} seconds') for command in commands: @@ -119,6 +162,19 @@ async def _run_commands(self, commands: Commands): await executor.run(command, is_api_instance=self.is_api_instance) async def run_command(self, command: Command) -> Result: + """Execute a single command and return its result. + + Looks up the appropriate executor for the command type and runs it. + One exception:``sftp`` commands are dispatched to the ``ssh`` executor. + + :param command: A command instance created via :meth:`Command.create`. + :returns: A :class:`Result` object containing ``stdout`` and ``returncode``. + Returns ``Result(None, None)`` if no executor is found. + + .. note:: + Commands running in backgound return + ``Result('Command started in background', 0)`` immediately. + """ command_type = 'ssh' if command.type == 'sftp' else command.type executor = self._get_executor(command_type) if executor: @@ -148,10 +204,13 @@ async def clean_session_stores(self): remote_executor.cleanup() async def main(self): - """The main function + """Execute the full playbook and clean up all sessions. - Passes the main playbook-commands to run_commands + Passes all commands from the playbook to :meth:`_run_commands`, then + tears down open sessions and background processes. Handles + :exc:`KeyboardInterrupt` gracefully. + :returns: ``0`` on completion. """ try: await self._run_commands(self.playbook.commands) diff --git a/src/attackmate/command.py b/src/attackmate/command.py index b2334e74..d851f0e4 100644 --- a/src/attackmate/command.py +++ b/src/attackmate/command.py @@ -3,12 +3,34 @@ class CommandRegistry: + """Registry mapping command types (and optional ``cmd`` values) to command classes. + + Command classes are registered via the :meth:`register` decorator. The registry + supports two lookup keys: + + - **type only** — e.g. ``"shell"`` maps to :class:`ShellCommand`. + - **type + cmd** — e.g. ``("sliver", "cd")`` maps to a more specific class. + + The more specific ``(type, cmd)`` key takes precedence over the ``type``-only key. + """ _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.""" + """Decorator to register a command class. + + :param type_: The command type string (e.g. ``"shell"``, ``"sleep"``). + :param cmd: An optional ``cmd`` value for more specific registration. + If provided, the class is registered under ``(type_, cmd)`` and will + only be returned when both match. This is a LEGACY pattern and should **NOT** be used anymore. + + Example:: + + @CommandRegistry.register("shell") + class ShellCommand(BaseCommand): + ... + """ def decorator(command_class: BaseCommand): if cmd: @@ -21,7 +43,16 @@ def decorator(command_class: BaseCommand): @classmethod def get_command_class(cls, type_: str, cmd: Optional[str] = None): - """Retrieve the command class based on type or type + cmd.""" + """Look up a registered command class. + + The ``(type_, cmd)`` key is checked first; if not found, falls back to + the ``type_``-only key. + + :param type_: The command type string. + :param cmd: An optional ``cmd`` value for specific lookup. + :returns: The registered command class. + :raises ValueError: If no class is registered for the given type and cmd. + """ if cmd and (type_, cmd) in cls._type_cmd_registry: return cls._type_cmd_registry[(type_, cmd)] if type_ in cls._type_registry: @@ -31,8 +62,31 @@ def get_command_class(cls, type_: str, cmd: Optional[str] = None): # Command interface class Command: + """Factory interface for creating command instances. + + Command classes are looked up from the :class:`CommandRegistry` based on + ``type`` and optionally ``cmd``, then instantiated with the remaining keyword + arguments. + """ @staticmethod def create(type: str, cmd: Optional[str] = None, **kwargs): - """return a command instance based on type and cmd.""" + """Create and return a command instance. + + Looks up the appropriate command class via :class:`CommandRegistry` and + instantiates it with the provided arguments. + + :param type: The command type (e.g. ``"shell"``, ``"sleep"``, ``"debug"``). + :param cmd: An optional ``cmd`` value used for more specific class lookup. + :param kwargs: Additional keyword arguments passed to the command class + constructor (e.g. ``seconds``, ``varstore``). + :returns: An instance of the resolved command class. + :raises ValueError: If no command class is registered for the given + ``type`` and ``cmd`` combination. + + Example:: + + command = Command.create(type="sleep", cmd="sleep", seconds="1") + command = Command.create(type="debug", cmd="$TEST", varstore=True) + """ CommandClass = CommandRegistry.get_command_class(type, cmd) return CommandClass(type=type, cmd=cmd, **kwargs) diff --git a/src/attackmate/executors/browser/sessionstore.py b/src/attackmate/executors/browser/sessionstore.py index 98945de6..d08fd2f9 100644 --- a/src/attackmate/executors/browser/sessionstore.py +++ b/src/attackmate/executors/browser/sessionstore.py @@ -109,9 +109,9 @@ def submit_command(self, cmd_name, *args, **kwargs): """ Called from outside to run a command in this thread. This is a synchronous call from the caller’s perspective: - - we place the command on cmd_queue, - - then wait for the response in res_queue - - re-raise the exception to preserve traceback if it failed + - we place the command on cmd_queue, + - then wait for the response in res_queue + - re-raise the exception to preserve traceback if it failed """ self.cmd_queue.put((cmd_name, args, kwargs)) result, error = self.res_queue.get() diff --git a/src/attackmate/executors/common/loopexecutor.py b/src/attackmate/executors/common/loopexecutor.py index 76cc011b..c6986195 100644 --- a/src/attackmate/executors/common/loopexecutor.py +++ b/src/attackmate/executors/common/loopexecutor.py @@ -51,7 +51,7 @@ def substitute_variables_in_command(self, command_obj, placeholders: dict): Replaces all occurrences of placeholders in *any* string attribute of the command_obj with the values from placeholders. E.g. if placeholders = {'LOOP_ITEM': 'https://example.com'}, - and command_obj.url = '$LOOP_ITEM', then it becomes 'https://example.com'. + and command_obj.url = '$LOOP_ITEM', then it becomes 'https://example.com'. """ for attr_name, attr_val in vars(command_obj).items(): if isinstance(attr_val, str) and '$' in attr_val: diff --git a/src/attackmate/variablestore.py b/src/attackmate/variablestore.py index 3f9c088a..4db56ab7 100644 --- a/src/attackmate/variablestore.py +++ b/src/attackmate/variablestore.py @@ -25,10 +25,34 @@ class ListTemplate(Template): class VariableStore: + """Stores and resolves variables for use in AttackMate playbooks. + + Variables are held in two separate namespaces: + + - **Scalar variables** — simple string key-value pairs, accessed as ``$varname``. + - **List variables** — ordered sequences of strings, accessed by index as + ``$varname[0]``, ``$varname[1]``, etc. + + All values are stored as plain Python ``str``, regardless of their original type. + Integer values are coerced to ``str`` on ingress. Variable substitution uses + Python's :class:`string.Template` syntax. + + If an environment variable prefixed with ``ATTACKMATE_`` exists with the same name + as a stored variable, it will override the stored value when + :meth:`replace_with_prefixed_env_vars` is called. + + Example:: + + store = VariableStore() + store.set_variable("HOST", "192.168.1.1") + store.substitute_str("Target: $HOST") # returns "Target: 192.168.1.1" + """ + def __init__(self): self.clear() def clear(self): + """Reset the store, removing all scalar variables and lists.""" self.lists: dict[str, list[str]] = {} self.variables: dict[str, str] = {} @@ -61,6 +85,14 @@ def get_lists_variables(self) -> dict[str, str]: return all_indexed_list_vars def from_dict(self, variables: Optional[dict]): + """Populate the store from a dictionary. + + Each key-value pair is passed to :meth:`set_variable`. Keys may optionally + include a leading ``$``, which is stripped before storing. + + :param variables: A dictionary of variable names to values. Non-dict values + are silently ignored. + """ if isinstance(variables, dict): for k, v in variables.items(): self.set_variable(k, v) @@ -86,6 +118,17 @@ def get_str(self, variable: str) -> str: raise VariableNotFound def substitute_str(self, template_str: str, blank: bool = False) -> str: + """Substitute all ``$variable`` references in a template string. + + Resolves both scalar variables and indexed list variables + (e.g. ``$mylist[0]``). + + :param template_str: A string containing ``$variable`` placeholders. + :param blank: If ``True``, unresolved references are replaced with ``''`` + and a :exc:`KeyError` causes the entire result to be ``''``. + If ``False`` (default), unresolved references are left as-is. + :returns: The string with all known variables substituted. + """ temp = ListTemplate(template_str) if blank: try: @@ -96,6 +139,16 @@ def substitute_str(self, template_str: str, blank: bool = False) -> str: return temp.safe_substitute(self.variables | self.get_lists_variables()) def set_variable(self, variable: str, value: str | list[str]): + """Store a variable by name. + + If ``value`` is an ``int`` it is coerced to ``str`` before storing. + If ``variable`` refers to a list index (e.g. ``mylist[0]``), the value + is written into the corresponding position of the existing list. + If ``value`` is a ``list``, it is stored as a list variable. + + :param variable: The variable name, with or without a leading ``$``. + :param value: The value to store. Integers are coerced to ``str``. + """ if isinstance(value, int): value = str(value) if isinstance(variable, str): @@ -110,6 +163,14 @@ def set_variable(self, variable: str, value: str | list[str]): self.lists[varname] = list(value) def get_variable(self, variable: str) -> str | list[str]: + """Retrieve a variable by name. + + Checks scalar variables first, then list variables. + + :param variable: The variable name (without ``$`` prefix). + :returns: The stored ``str`` value or ``list[str]``. + :raises VariableNotFound: If no variable with that name exists. + """ if variable in self.variables: return self.variables[variable] if variable in self.lists: @@ -117,17 +178,38 @@ def get_variable(self, variable: str) -> str | list[str]: raise VariableNotFound def substitute(self, data: Any, blank: bool = False) -> Any: + """Substitute variables in ``data`` if it is a string. + + Non-string values are returned unchanged. + + :param data: The value to substitute into. Only ``str`` values are processed. + :param blank: If ``True``, unresolved variable references are replaced with + an empty string. If ``False`` (default), they are left as-is. + :returns: The substituted string, or the original value if not a string. + """ if isinstance(data, str): return self.substitute_str(data, blank) else: return data def get_prefixed_env_vars(self, prefix: str = 'ATTACKMATE_') -> dict[str, str]: + """Return all environment variables that start with ``prefix``, stripped of the prefix. + + :param prefix: The prefix to filter and strip. Defaults to ``'ATTACKMATE_'``. + :returns: A dictionary mapping unprefixed names to their environment values. + """ prefixed_env_vars = {k[len(prefix):]: v for k, v in os.environ.items() if k.startswith(prefix)} return prefixed_env_vars def replace_with_prefixed_env_vars(self): - """Replaces the current variables with corresponding prefixed environment variables if they exist.""" + """Override stored variables with matching prefixed environment variables. + + For each scalar variable currently in the store, if an environment variable + named ``ATTACKMATE_`` exists, its value replaces the stored one. + + Example: the stored variable ``FOO`` is overridden by the environment + variable ``ATTACKMATE_FOO`` if it is set. + """ env_vars = self.get_prefixed_env_vars() for var_name in list(self.variables.keys()):