From d62366da454ec6bd48427d1a91de027241549a6c Mon Sep 17 00:00:00 2001
From: Caleb Stewart <caleb.stewart94@gmail.com>
Date: Sun, 17 May 2020 02:29:51 -0400
Subject: [PATCH] Run/local command and shortcuts

Added the "run" and "local" commands for remote and local command
execution respectively and the "shortcut" command to allow for shortcuts
like "!ls" for local commands and "@ls" for remote commands.
---
 data/pwncatrc               |  8 +++++++-
 pwncat/commands/__init__.py | 34 +++++++++++++++++++++-------------
 pwncat/commands/local.py    | 19 +++++++++++++++++++
 pwncat/commands/run.py      | 18 ++++++++++++++++++
 pwncat/commands/shortcut.py | 24 ++++++++++++++++++++++++
 5 files changed, 89 insertions(+), 14 deletions(-)
 create mode 100644 pwncat/commands/local.py
 create mode 100644 pwncat/commands/run.py
 create mode 100644 pwncat/commands/shortcut.py

diff --git a/data/pwncatrc b/data/pwncatrc
index 2dd945b..8652340 100644
--- a/data/pwncatrc
+++ b/data/pwncatrc
@@ -17,6 +17,12 @@ set on_load {
 bind s "sync"
 bind c "set state command"
 
-# Create shortcut aliases for commands
+# Create aliases for commands
 alias up upload
 alias down download
+
+# Shortcuts allow single-character prefix which indicate the entire command
+# string be passed as the arguments to a specific command. For example:
+# "!ls" run "local ls" given the below directives
+shortcut ! local
+shortcut @ run
diff --git a/pwncat/commands/__init__.py b/pwncat/commands/__init__.py
index c93b936..880d7a2 100644
--- a/pwncat/commands/__init__.py
+++ b/pwncat/commands/__init__.py
@@ -146,6 +146,7 @@ class CommandParser:
 
         self.loading_complete = False
         self.aliases: Dict[str, CommandDefinition] = {}
+        self.shortcuts: Dict[str, CommandDefinition] = {}
 
     @property
     def loaded(self):
@@ -220,24 +221,31 @@ class CommandParser:
             util.error(e.args[0])
             return
 
-        # Search for a matching command
-        for command in self.commands:
-            if command.PROG == argv[0]:
-                break
+        if argv[0][0] in self.shortcuts:
+            command = self.shortcuts[argv[0][0]]
+            argv[0] = argv[0][1:]
+            args = argv
         else:
-            if argv[0] in self.aliases:
-                command = self.aliases[argv[0]]
+            # Search for a matching command
+            for command in self.commands:
+                if command.PROG == argv[0]:
+                    break
             else:
-                util.error(f"{argv[0]}: unknown command")
+                if argv[0] in self.aliases:
+                    command = self.aliases[argv[0]]
+                else:
+                    util.error(f"{argv[0]}: unknown command")
+                    return
+
+            if not self.loading_complete and not command.LOCAL:
+                util.error(
+                    f"{argv[0]}: non-local commands cannot run until after session setup."
+                )
                 return
 
-        if not self.loading_complete and not command.LOCAL:
-            util.error(
-                f"{argv[0]}: non-local commands cannot run until after session setup."
-            )
-            return
+            args = argv[1:]
 
-        args = [a.encode("utf-8").decode("unicode_escape") for a in argv[1:]]
+        args = [a.encode("utf-8").decode("unicode_escape") for a in args]
 
         try:
             # Parse the arguments
diff --git a/pwncat/commands/local.py b/pwncat/commands/local.py
new file mode 100644
index 0000000..30f8a9c
--- /dev/null
+++ b/pwncat/commands/local.py
@@ -0,0 +1,19 @@
+#!/usr/bin/env python3
+import subprocess
+
+from pwncat.commands import CommandDefinition, Complete
+from pwncat.commands.base import parameter
+
+
+class Command(CommandDefinition):
+
+    PROG = "local"
+    ARGS = {
+        "argv": parameter(
+            Complete.NONE, nargs="+", help="the local shell command to run"
+        )
+    }
+    LOCAL = True
+
+    def run(self, args):
+        subprocess.run(args.argv, shell=True)
diff --git a/pwncat/commands/run.py b/pwncat/commands/run.py
new file mode 100644
index 0000000..0db4b2e
--- /dev/null
+++ b/pwncat/commands/run.py
@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+import sys
+
+import pwncat
+from pwncat.commands.base import CommandDefinition, Complete, parameter
+
+
+class Command(CommandDefinition):
+
+    PROG = "run"
+    ARGS = {
+        "argv": parameter(
+            Complete.NONE, nargs="+", help="The command to run on the remote host"
+        )
+    }
+
+    def run(self, args):
+        sys.stdout.buffer.write(pwncat.victim.run(args.argv))
diff --git a/pwncat/commands/shortcut.py b/pwncat/commands/shortcut.py
new file mode 100644
index 0000000..037ecd4
--- /dev/null
+++ b/pwncat/commands/shortcut.py
@@ -0,0 +1,24 @@
+import pwncat
+from pwncat.commands import CommandDefinition
+from pwncat.commands.base import parameter, Complete
+
+
+class Command(CommandDefinition):
+
+    PROG = "shortcut"
+    ARGS = {
+        "prefix": parameter(
+            Complete.NONE, help="the prefix character used for the shortcut"
+        ),
+        "command": parameter(Complete.NONE, help="the command to execute"),
+    }
+    LOCAL = True
+
+    def run(self, args):
+
+        for command in pwncat.victim.command_parser.commands:
+            if command.PROG == args.command:
+                pwncat.victim.command_parser.shortcuts[args.prefix] = command
+                return
+
+        self.parser.error(f"{args.command}: no such command")