Added bind and alias commands to fully control configuration through command scripting.

data/pwncatrc

@@ -9,13 +9,27 @@ set backdoor_user "pwncat"
 set backdoor_pass "pwncat"
 
 # This will fail because we haven't finished loading
-busybox -i
 
 set on_load {
 	# This will succeed because `on_load` runs after the session is established.
-	busybox -i
 }
 
+# Examples of command bindings
+bind s "sync"
+bind c "set state command"
+bind b {
+	busybox --install
+	busybox --status
+}
+bind p {
+	privesc --list
+}
+
+# Create shortcut aliases for commands
+alias up upload
+alias down download
+alias down
+
 # This would run at start after a successful connection
 # privesc -l
 

pwncat/commands/__init__.py

@@ -21,6 +21,7 @@ from pygments.styles import get_style_by_name
 from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
 from prompt_toolkit.history import InMemoryHistory
 from typing import Dict, Any, List, Iterable
+from colorama import Fore
 from enum import Enum, auto
 import argparse
 import pkgutil
@@ -146,6 +147,7 @@ class CommandParser:
 
         self.pty = pty
         self.loading_complete = False
+        self.aliases: Dict[str, CommandDefinition] = {}
 
     @property
     def loaded(self):
@@ -170,7 +172,9 @@ class CommandParser:
             try:
                 self.dispatch_line(command)
             except Exception as exc:
-                util.error(f"{name}: command: {str(exc)}")
+                util.error(
+                    f"{Fore.CYAN}{name}{Fore.RESET}: {Fore.YELLOW}{command}{Fore.RESET}: {str(exc)}"
+                )
                 break
 
     def run_single(self):
@@ -222,6 +226,9 @@ class CommandParser:
         for command in self.commands:
             if command.PROG == argv[0]:
                 break
+        else:
+            if argv[0] in self.aliases:
+                command = self.aliases[argv[0]]
             else:
                 util.error(f"{argv[0]}: unknown command")
                 return

pwncat/commands/alias.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+from pwncat.commands.base import CommandDefinition, Complete, parameter
+from colorama import Fore
+
+
+class Command(CommandDefinition):
+    """ Alias an existing command with a new name. Specifying no alias or command
+    will list all aliases. Specifying an alias with no command will remove the 
+    alias if it exists. """
+
+    def get_command_names(self):
+        return [c.PROG for c in self.cmdparser.commands]
+
+    PROG = "alias"
+    ARGS = {
+        "alias": parameter(Complete.NONE, help="name for the new alias", nargs="?"),
+        "command": parameter(
+            Complete.CHOICES,
+            metavar="COMMAND",
+            choices=get_command_names,
+            help="the command the new alias will use",
+            nargs="?",
+        ),
+    }
+    LOCAL = True
+
+    def run(self, args):
+        if args.alias is None:
+            for name, command in self.cmdparser.aliases.items():
+                print(
+                    f" {Fore.CYAN}{name}{Fore.RESET} \u2192 "
+                    f"{Fore.YELLOW}{command.PROG}{Fore.RESET}"
+                )
+        elif args.command is not None:
+            # This is safe because of "choices" in the argparser
+            self.cmdparser.aliases[args.alias] = [
+                c for c in self.cmdparser.commands if c.PROG == args.command
+            ][0]
+        else:
+            del self.cmdparser.aliases[args.alias]

pwncat/commands/bind.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+from prompt_toolkit.input.ansi_escape_sequences import REVERSE_ANSI_SEQUENCES
+from prompt_toolkit.keys import ALL_KEYS, Keys
+from pwncat.commands.base import CommandDefinition, Complete, parameter
+from pwncat.config import KeyType
+from pwncat import util
+from colorama import Fore
+import string
+
+
+class Command(CommandDefinition):
+
+    PROG = "bind"
+    ARGS = {
+        "key": parameter(
+            Complete.NONE,
+            metavar="KEY",
+            type=KeyType,
+            help="The key to map after your prefix",
+            nargs="?",
+        ),
+        "script": parameter(
+            Complete.NONE, help="The script to run when the key is pressed", nargs="?",
+        ),
+    }
+    LOCAL = True
+
+    def run(self, args):
+        if args.key is None:
+            util.info("currently assigned key-bindings:")
+            for key, binding in self.pty.config.bindings.items():
+                print(
+                    f" {Fore.CYAN}{key}{Fore.RESET} = {Fore.YELLOW}{repr(binding)}{Fore.RESET}"
+                )
+        elif args.key is not None and args.script is None:
+            if args.key in self.pty.config.bindings:
+                del self.pty.config.bindings[args.key]
+        else:
+            self.pty.config.bindings[args.key] = args.script

pwncat/commands/set.py

@@ -8,7 +8,7 @@ class Command(CommandDefinition):
     """ Set variable runtime variable parameters for pwncat """
 
     def get_config_variables(self):
-        return list(self.pty.config.values)
+        return ["state"] + list(self.pty.config.values)
 
     PROG = "set"
     ARGS = {
@@ -26,7 +26,16 @@ class Command(CommandDefinition):
     LOCAL = True
 
     def run(self, args):
-        if args.variable is not None and args.value is not None:
+        if (
+            args.variable is not None
+            and args.variable == "state"
+            and args.value is not None
+        ):
+            try:
+                self.pty.state = util.State._member_map_[args.value.upper()]
+            except KeyError:
+                util.error(f"{args.value}: invalid state")
+        elif args.variable is not None and args.value is not None:
             try:
                 self.pty.config[args.variable] = args.value
             except ValueError as exc:

pwncat/config.py

@@ -1,7 +1,10 @@
 #!/usr/bin/env python3
-from prompt_toolkit.input.ansi_escape_sequences import REVERSE_ANSI_SEQUENCES
+from prompt_toolkit.input.ansi_escape_sequences import (
+    REVERSE_ANSI_SEQUENCES,
+    ANSI_SEQUENCES,
+)
 from prompt_toolkit.keys import ALL_KEYS, Keys
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Union
 import commentjson as json
 import ipaddress
 import re
@@ -50,8 +53,11 @@ class Config:
 
         # Basic key-value store w/ typing
         self.values: Dict[str, Dict[str, Any]] = {
-            "lhost": {"value": None, "type": ipaddress.ip_address},
-            "prefix": {"value": "C-k", "type": KeyType},
+            "lhost": {
+                "value": ipaddress.ip_address("127.0.0.1"),
+                "type": ipaddress.ip_address,
+            },
+            "prefix": {"value": KeyType("c-k"), "type": KeyType},
             "privkey": {"value": "data/pwncat", "type": local_file_type},
             "backdoor_user": {"value": "pwncat", "type": str},
             "backdoor_pass": {"value": "pwncat", "type": str},
@@ -60,7 +66,25 @@ class Config:
 
         # Map ascii escape sequences or printable bytes to lists of commands to
         # run.
-        self.bindings: Dict[bytes, str] = {}
+        self.bindings: Dict[KeyType, str] = {
+            KeyType("c-d"): "pass",
+            KeyType("s"): "sync",
+            KeyType("c"): "set state command",
+        }
+
+    def binding(self, name_or_value: Union[str, bytes]) -> str:
+        """ Get a key binding by it's key name or key value. """
+
+        if isinstance(name_or_value, bytes):
+            binding = [
+                b for key, b in self.bindings.items() if key.value == name_or_value
+            ]
+            if not binding:
+                raise KeyError("no such key binding")
+            return binding[0]
+
+        key = KeyType(name_or_value)
+        return self.bindings[key]
 
     def __getitem__(self, name: str) -> Any:
         """ Get a configuration item """

pwncat/pty.py

@@ -38,7 +38,7 @@ from pwncat.file import RemoteBinaryPipe
 from pwncat.lexer import LocalCommandLexer, PwncatStyle
 from pwncat.gtfobins import GTFOBins, Capability, Stream
 from pwncat.commands import CommandParser
-from pwncat.config import Config
+from pwncat.config import Config, KeyType
 
 from colorama import Fore
 
@@ -216,6 +216,7 @@ class PtyHandler:
         }
         self.gtfo: GTFOBins = GTFOBins("data/gtfobins.json", self.which)
         self.default_privkey = "./data/pwncat"
+        self.has_prefix = False
         self.command_parser = CommandParser(self)
 
         # Run the configuration script
@@ -549,29 +550,41 @@ class PtyHandler:
         r""" Process a new byte of input from stdin. This is to catch "\r~C" and open
         a local prompt """
 
-        if self.input == b"":
-            # Enter commmand mode after C-d
-            if data == b"\x04":
-                # Clear line
-                self.client.send(b"\x15")
-                # Enter command mode
+        if self.has_prefix:
+            if data == self.config["prefix"].value:
+                self.client.send(data)
+            else:
+                try:
+                    binding = self.config.binding(data)
+
+                    # Pass is a special case that can be used at the beginning of a
+                    # command.
+                    if binding.strip().startswith("pass"):
+                        self.client.send(data)
+                        binding = binding.lstrip("pass")
+
+                    self.restore_local_term()
+                    sys.stdout.write("\n")
+
+                    # Evaluate the script
+                    self.command_parser.eval(binding, "<binding>")
+
+                    self.flush_output()
+                    self.client.send(b"\n")
+                    self.saved_term_state = util.enter_raw_mode()
+
+                except KeyError:
+                    pass
+            self.has_prefix = False
+        elif data == self.config["prefix"].value:
+            self.has_prefix = True
+        elif data == KeyType("c-d").value:
+            # Don't allow exiting the remote prompt with C-d
+            # you should have a keybinding for "<prefix> C-d" to actually send
+            # C-d.
             self.state = State.COMMAND
-            # C-k is the prefix character
-            elif data == b"\x0b":
-                self.input = data
         else:
             self.client.send(data)
-        else:
-            # "C-k c" to enter command mode
-            if data == b"c":
-                self.client.send(b"\x15")
-                self.state = State.SINGLE
-            elif data == b":":
-                self.state = State.SINGLE
-            # "C-k C-k" or "C-k C-d" sends the second byte
-            elif data == b"\x0b" or data == b"\x04":
-                self.client.send(data)
-            self.input = b""
 
     def recv(self) -> bytes:
         """ Recieve data from the client """