Added subprocess option to get file-like access to command output, and a downloader that reuses the open socket connection for fast downloads

pwncat/downloader/__init__.py

@@ -7,15 +7,17 @@ from pwncat.downloader.nc import NetcatDownloader
 from pwncat.downloader.curl import CurlDownloader
 from pwncat.downloader.shell import ShellDownloader
 from pwncat.downloader.bashtcp import BashTCPDownloader
+from pwncat.downloader.raw import RawShellDownloader
 
 all_downloaders = [
     NetcatDownloader,
     CurlDownloader,
     ShellDownloader,
     BashTCPDownloader,
+    RawShellDownloader,
 ]
 downloaders = [NetcatDownloader, CurlDownloader]
-fallback = ShellDownloader
+fallback = RawShellDownloader
 
 
 def get_names() -> List[str]:

pwncat/downloader/base.py

@@ -23,7 +23,12 @@ class Downloader:
     def check(cls, pty: "pwncat.pty.PtyHandler") -> bool:
         """ Check if the given PTY connection can support this downloader """
         for binary in cls.BINARIES:
-            if pty.which(binary) is None:
+            if isinstance(binary, list) or isinstance(binary, tuple):
+                for equivalent in binary:
+                    if pty.which(equivalent):
+                        return
+            elif pty.which(binary) is not None:
+                return
             raise DownloadError(f"required remote binary not found: {binary}")
 
     def __init__(self, pty: "pwncat.pty.PtyHandler", remote_path: str, local_path: str):

pwncat/downloader/raw.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+from typing import Generator, Callable
+from io import BufferedReader
+import base64
+import shlex
+
+from pwncat.downloader.base import Downloader, DownloadError
+from pwncat import util
+
+
+class RawShellDownloader(Downloader):
+
+    NAME = "raw"
+    BINARIES = [("dd", "cat")]
+    BLOCKSZ = 8192
+
+    def command(self) -> Generator[str, None, None]:
+        """ Yield list of commands to transfer the file """
+
+        remote_path = shlex.quote(self.remote_path)
+        blocksz = 1024 * 1024
+        binary = self.pty.which("dd")
+
+        if binary is None:
+            binary = self.pty.which("cat")
+
+        if "dd" in binary:
+            pipe = self.pty.subprocess(f"dd if={remote_path} bs={blocksz} 2>/dev/null")
+        else:
+            pipe = self.pty.subprocess(f"cat {remote_path}")
+
+        try:
+            with open(self.local_path, "wb") as filp:
+                util.copyfileobj(pipe, filp, self.on_progress)
+        finally:
+            self.on_progress(0, -1)
+            pipe.close()
+
+        return False
+
+    def serve(self, on_progress: Callable):
+        """ We don't need to start a server, but we do need to save the
+            callback """
+        self.on_progress = on_progress

pwncat/file.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+from io import RawIOBase
+import socket
+
+
+class RemoteBinaryPipe(RawIOBase):
+    """ Encapsulate a piped interaction with a remote process. The remote PTY
+    should have been placed in raw mode prior to this object being created, and 
+    the appropriate flags in pty already modified. If EOF is found or the object
+    is closed, it will restore the state of the terminal (w/ `reset`). No further
+    reading or writing will be allowed. """
+
+    def __init__(self, pty: "pwncat.pty.PtyHandler", delim: bytes, binary: bool):
+        self.pty = pty
+        self.delim = delim
+        self.eof = 0
+        self.next_eof = False
+        self.binary = binary
+        self.split_eof = b""
+
+    def on_eof(self):
+        if self.eof:
+            return
+
+        # Set eof flag
+        self.eof = 1
+
+        if self.binary:
+            # Reset the terminal
+            self.pty.reset()
+            # Send a bare echo, and read all data to ensure we don't clobber the
+            # output of the user's terminal
+            self.pty.run("echo")
+
+    def close(self):
+        if self.eof:
+            return
+
+        # Kill the last job. This should be us.
+        self.pty.run("kill -9 %%", wait=False)
+
+        # Cleanup
+        self.on_eof()
+
+    # def read(self, size: int = -1):
+    #     if self.eof == -1:
+    #         self.on_eof()
+
+    #     if self.eof:
+    #         return b""
+
+    #     if size == -1:
+    #         data = b""
+    #         while self.delim not in data:
+    #             data += self.pty.client.recv(1024 * 1024)
+    #         data = data.split(self.delim)[0]
+    #         self.eof = -1
+    #     else:
+    #         data = self.pty.client.recv(size)
+    #         if self.delim in data:
+    #             self.eof = -1
+    #             data = data.split(self.delim)[0]
+
+    #     return data
+
+    def readinto(self, b: bytearray):
+        if self.eof:
+            return 0
+
+        if isinstance(b, memoryview):
+            obj = b.obj
+        else:
+            obj = b
+
+        # Receive the data
+        n = self.pty.client.recv_into(b)
+
+        # Check for EOF
+        if self.delim in obj:
+            self.on_eof()
+            n = obj.find(self.delim)
+            return n
+        else:
+            # Check for EOF split across blocks
+            for i in range(1, len(self.delim) - 1):
+                # See if a piece of the delimeter is at the end of this block
+                piece = self.delim[:-i]
+                if bytes(b[-len(piece) :]) == piece:
+                    try:
+                        # Peak the next bytes, to see if this is actually the
+                        # delimeter
+                        rest = self.pty.client.recv(
+                            i, socket.MSG_PEEK | socket.MSG_DONTWAIT
+                        )
+                    except (socket.error, BlockingIOError):
+                        rest = b""
+                    # It is!
+                    if (piece + rest) == self.delim:
+                        # Receive the delimeter
+                        self.pty.client.recv(i)
+                        # Adjust result
+                        n -= len(piece)
+                        # Set EOF for next read
+                        self.on_eof()
+
+        return n
+
+    def flush_read(self):
+        """ read all until eof and ignore it """
+        for block in iter(lambda: self.read(1024 * 1024), b""):
+            pass
+
+    def write(self, data: bytes):
+        return self.pty.client.send(data)

pwncat/pty.py

@@ -28,6 +28,7 @@ import os
 from pwncat import util
 from pwncat import downloader, uploader, privesc
 from pwncat.lexer import LocalCommandLexer
+from pwncat.file import RemoteBinaryPipe
 from colorama import Fore
 
 
@@ -260,7 +261,9 @@ class PtyHandler:
         self.run("unset HISTFILE; export HISTCONTROL=ignorespace")
 
         util.info("setting terminal prompt", overlay=True)
-        self.run(f'export PS1="{self.remote_prefix} $PS1"')
+        self.run(
+            f'export SAVED_PS1="$PS1"; export PS1="{self.remote_prefix} $SAVED_PS1"'
+        )
 
         # Locate interesting binaries
         # The auto-resolving doesn't work correctly until we have a pty
@@ -303,7 +306,9 @@ class PtyHandler:
         self.has_prompt = True
 
         util.info("setting terminal prompt", overlay=True)
-        self.run(f'export PS1="{self.remote_prefix} $PS1"')
+        self.run(
+            f'export SAVED_PS1="$PS1"; export PS1="{self.remote_prefix} $SAVED_PS1"'
+        )
 
         # Make sure HISTFILE is unset in this PTY (it resets when a pty is
         # opened)
@@ -452,6 +457,9 @@ class PtyHandler:
                         result = self.run(line[1:])
                         sys.stdout.buffer.write(result)
                         continue
+                    elif line[0] == "-":
+                        self.run(line[1:], wait=False)
+                        continue
 
                 argv = shlex.split(line)
 
@@ -539,6 +547,12 @@ class PtyHandler:
 
             def on_progress(copied, blocksz):
                 """ Update the progress bar """
+                if blocksz == -1:
+                    counter.stopped = True
+                    counter.done = True
+                    pb.invalidate()
+                    return
+
                 counter.items_completed += blocksz
                 if counter.items_completed >= counter.total:
                     counter.done = True
@@ -546,15 +560,11 @@ class PtyHandler:
                 if (time.time() - last_update) > 0.1:
                     pb.invalidate()
 
-            download.serve(on_progress)
-
-            download.command()
-
             try:
+                download.serve(on_progress)
+                if download.command():
                     while not counter.done:
-                    time.sleep(0.1)
-            except KeyboardInterrupt:
-                pass
+                        time.sleep(0.2)
             finally:
                 download.shutdown()
 
@@ -653,6 +663,11 @@ class PtyHandler:
             help_msg = getattr(self, c).__doc__
             print(f"{c[3:]:15s}{help_msg}")
 
+    def do_reset(self, argv):
+        """ Reset the remote terminal (calls sync, reset, and sets PS1) """
+        self.reset()
+        self.do_sync([])
+
     def run(self, cmd, wait=True) -> bytes:
         """ Run a command in the context of the remote host and return the
         output. This is run synchrounously.
@@ -689,7 +704,7 @@ class PtyHandler:
         if delim:
             command = f" echo _PWNCAT_STARTDELIM_; {cmd}; echo _PWNCAT_ENDDELIM_"
         else:
-            command = cmd
+            command = f" {cmd}"
 
         response = b""
 
@@ -707,9 +722,80 @@ class PtyHandler:
 
         return b"_PWNCAT_ENDDELIM_"
 
+    def subprocess(self, cmd) -> RemoteBinaryPipe:
+        """ Create an asynchronous child on the remote end and return a
+        file-like object which can communicate with it's standard output. The 
+        remote terminal is placed in raw mode with no-echo first, and the
+        command is run on a separate background shell w/ no standard input. The
+        output of the command can be retrieved through the returned file-like
+        object. You **must** either call `close()` of the pipe, or read until
+        eof, or the PTY will not be restored to a normal state.
+
+        If `close()` is called prior to EOF, the remote process will be killed,
+        and any remaining output will be flushed prior to resetting the terminal.
+        """
+
+        if isinstance(cmd, list):
+            cmd = shlex.join(cmd)
+
+        sdelim = "_PWNCAT_STARTDELIM_"
+        edelim = "_PWNCAT_ENDDELIM_"
+
+        # List of ";" separated commands that will be run
+        command = []
+        # Clear the prompt, or it will get displayed in our output due to the
+        # background task
+        command.append("export PS1=")
+        # Needed to disable job control messages in bash
+        command.append("set +m")
+        # This is gross, but it allows us to recieve stderr and stdout, while
+        # ignoring the job control start message.
+        command.append(
+            f"{{ echo {sdelim}; {cmd} && echo {edelim} || echo {edelim} 2>&1 & }} 2>/dev/null"
+        )
+        # Re-enable normal job control in bash
+        command.append("set -m")
+
+        # Join them all into one command
+        command = ";".join(command).encode("utf-8")
+
+        # Enter raw mode w/ no echo on the remote terminal
+        # DANGER
+        self.raw(echo=False)
+
+        self.client.sendall(command + b"\n")
+        self.recvuntil(sdelim)
+        self.recvuntil("\n")
+
+        # Bash sends some bullshit that messes up terminals. Check to see if it
+        # is there, and ignore it if it is
+        try:
+            data = self.client.recv(2, socket.MSG_PEEK | socket.MSG_DONTWAIT)
+        except (BlockingIOError, socket.error):
+            data = b""
+
+        if data == b"\x1b_":
+            self.recvuntil(b"\x1b\\")
+
+        return RemoteBinaryPipe(self, edelim.encode("utf-8"), True)
+
+    def raw(self, echo: bool = False):
+        self.run("stty raw -echo", wait=False)
+        self.has_cr = False
+        self.has_echo = False
+
+    def reset(self):
+        self.run("reset", wait=False)
+        self.has_cr = True
+        self.has_echo = True
+        self.run(f'export PS1="{self.remote_prefix} $SAVED_PS1"')
+
     def recvuntil(self, needle: bytes, flags=0):
         """ Recieve data from the client until the specified string appears """
 
+        if isinstance(needle, str):
+            needle = needle.encode("utf-8")
+
         result = b""
         while not result.endswith(needle):
             try: