From b29dcde25234f68eb1b9bb7b54a4b1a83de6b307 Mon Sep 17 00:00:00 2001
From: John Hammond <johnhammond010@gmail.com>
Date: Thu, 7 May 2020 21:46:47 -0400
Subject: [PATCH 1/2] Added support for a remote command to be run in the local
 context with at @ prefix

---
 pwncat/pty.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/pwncat/pty.py b/pwncat/pty.py
index 99642ff..d103204 100644
--- a/pwncat/pty.py
+++ b/pwncat/pty.py
@@ -276,10 +276,15 @@ class PtyHandler:
                     self.enter_raw()
                     continue
 
-                if len(line) > 0 and line[0] == "!":
-                    # Allow running shell commands
-                    subprocess.run(line[1:], shell=True)
-                    continue
+                if len(line) > 0:
+                    if line[0] == "!":
+                        # Allow running shell commands
+                        subprocess.run(line[1:], shell=True)
+                        continue
+                    elif line[0] == "@":
+                        result = self.run(line[1:])
+                        sys.stdout.buffer.write(result)
+                        continue
 
                 argv = shlex.split(line)
 
@@ -351,7 +356,7 @@ class PtyHandler:
         with ProgressBar(
             [("#888888", "downloading with "), ("fg:ansiyellow", f"{download.NAME}")]
         ) as pb:
-            counter = pb(range(os.path.getsize(path)))
+            counter = pb(range(size))
             last_update = time.time()
 
             def on_progress(copied, blocksz):

From 283d313abe8067ce7b6c6af5a098df02dc9f1b99 Mon Sep 17 00:00:00 2001
From: John Hammond <johnhammond010@gmail.com>
Date: Thu, 7 May 2020 22:08:17 -0400
Subject: [PATCH 2/2] Changed colors for prompts and corrected sync command to
 run on all one line

---
 pwncat/pty.py  | 25 ++++++++++++++-----------
 pwncat/util.py |  2 +-
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/pwncat/pty.py b/pwncat/pty.py
index d103204..17f5a60 100644
--- a/pwncat/pty.py
+++ b/pwncat/pty.py
@@ -63,9 +63,14 @@ class PtyHandler:
         self.lhost = None
         self.known_binaries = {}
         self.vars = {"lhost": util.get_ip_addr()}
+        self.remote_prefix = "\\[\\033[01;31m\\](remote)\\033[00m\\]"
         self.remote_prompt = "\\[\\033[01;32m\\]\\u@\\h\\[\\033[00m\\]:\\[\\033[01;34m\\]\\w\\[\\033[00m\\]\\$"
         self.prompt = PromptSession(
-            [("", "(local) "), ("#ff0000", "pwncat"), ("", "$ ")]
+            [
+                ("fg:ansiyellow bold", "(local) "),
+                ("fg:ansimagenta bold", "pwncat"),
+                ("", "$ "),
+            ]
         )
         self.binary_aliases = {
             "python": [
@@ -81,9 +86,9 @@ class PtyHandler:
         }
 
         # We should always get a response within 3 seconds...
-        self.client.settimeout(3)
+        self.client.settimeout(1)
 
-        util.info("probing for prompt...", overlay=False)
+        util.info("probing for prompt...", overlay=True)
         start = time.time()
         prompt = b""
         try:
@@ -131,10 +136,10 @@ class PtyHandler:
 
         # Ensure history is disabled
         util.info("disabling remote command history", overlay=True)
-        self.run("unset HISTFILE")
+        self.run("unset HISTFILE; export HISTCONTROL=ignorespace")
 
         util.info("setting terminal prompt", overlay=True)
-        self.run(f'export PS1="(remote) {self.remote_prompt} "')
+        self.run(f'export PS1="{self.remote_prefix} $PS1"')
 
         # Locate interesting binaries
         # The auto-resolving doesn't work correctly until we have a pty
@@ -177,11 +182,11 @@ class PtyHandler:
         self.has_prompt = True
 
         util.info("setting terminal prompt", overlay=True)
-        self.run(f'export PS1="(remote) {self.remote_prompt} "')
+        self.run(f'export PS1="{self.remote_prefix} $PS1"')
 
         # Make sure HISTFILE is unset in this PTY (it resets when a pty is
         # opened)
-        self.run("unset HISTFILE")
+        self.run("unset HISTFILE; export HISTCONTROL=ignorespace")
 
         # Synchronize the terminals
         util.info("synchronizing terminal state", overlay=True)
@@ -462,9 +467,7 @@ class PtyHandler:
         TERM = os.environ.get("TERM", "xterm")
         columns, rows = os.get_terminal_size(0)
 
-        self.run(f"stty rows {rows}")
-        self.run(f"stty columns {columns}")
-        self.run(f'export TERM="{TERM}"')
+        self.run(f"stty rows {rows}; stty columns {columns}; export TERM='{TERM}'")
 
     def do_set(self, argv):
         """ Set or view the currently assigned variables """
@@ -512,7 +515,7 @@ class PtyHandler:
         EOL = b"\r" if has_pty else b"\n"
 
         if wait:
-            command = f"echo _PWNCAT_DELIM_; {cmd}; echo _PWNCAT_DELIM_"
+            command = f" echo _PWNCAT_DELIM_; {cmd}; echo _PWNCAT_DELIM_"
         else:
             command = cmd
 
diff --git a/pwncat/util.py b/pwncat/util.py
index ca37566..8c56383 100644
--- a/pwncat/util.py
+++ b/pwncat/util.py
@@ -56,7 +56,7 @@ def enter_raw_mode():
     """
 
     info("setting terminal to raw mode and disabling echo", overlay=True)
-    success("pwncat is ready\n", overlay=True)
+    success("pwncat is ready 🐈\n", overlay=True)
 
     # Ensure we don't have any weird buffering issues
     sys.stdout.flush()