Merge branch 'master' of https://github.com/calebstewart/pwncat

pwncat/commands/__init__.py

@@ -261,7 +261,9 @@ class CommandParser:
             command = self.shortcuts[argv[0][0]]
             argv[0] = argv[0][1:]
             args = argv
+            line = line[1:]
         else:
+            line = f"{argv[0]} ".join(line.split(f"{argv[0]} ")[1:])
             # Search for a matching command
             for command in self.commands:
                 if command.PROG == argv[0]:
@@ -290,7 +292,10 @@ class CommandParser:
                 prog_name = temp_name
 
             # Parse the arguments
-            args = command.parser.parse_args(args)
+            if command.parser:
+                args = command.parser.parse_args(args)
+            else:
+                args = line
 
             # Run the command
             command.run(args)
@@ -315,17 +320,18 @@ class CommandLexer(RegexLexer):
         for command in commands:
             root.append(("^" + re.escape(command.PROG), Name.Function, command.PROG))
             mode = []
-            for args, descr in command.ARGS.items():
-                for arg in args.split(","):
-                    if not arg.startswith("-"):
-                        continue
-                    if descr[0] != Complete.NONE:
-                        # Enter param state
-                        mode.append((r"\s+" + re.escape(arg), descr[1], "param"))
-                    else:
-                        # Don't enter param state
-                        mode.append((r"\s+" + re.escape(arg), descr[1]))
-            mode.append((r"\s+(\-\-help|\-h)", Name.Label))
+            if command.ARGS is not None:
+                for args, descr in command.ARGS.items():
+                    for arg in args.split(","):
+                        if not arg.startswith("-"):
+                            continue
+                        if descr[0] != Complete.NONE:
+                            # Enter param state
+                            mode.append((r"\s+" + re.escape(arg), descr[1], "param"))
+                        else:
+                            # Don't enter param state
+                            mode.append((r"\s+" + re.escape(arg), descr[1]))
+                mode.append((r"\s+(\-\-help|\-h)", Name.Label))
             mode.append((r"\"", String, "string"))
             mode.append((r".", Text))
             cls.tokens[command.PROG] = mode
@@ -408,25 +414,26 @@ class CommandCompleter(Completer):
         for command in commands:
             self.layers[command.PROG] = [None, [], {}]
             option_names = []
-            for name_list, descr in command.ARGS.items():
-                name_list = name_list.split(",")
-                if descr[0] == Complete.CHOICES:
-                    completer = WordCompleter(descr[3]["choices"])
-                elif descr[0] == Complete.LOCAL_FILE:
-                    completer = local_file_completer
-                elif descr[0] == Complete.REMOTE_FILE:
-                    completer = remote_file_completer
-                elif descr[0] == Complete.NONE:
-                    completer = None
-                if len(name_list) == 1 and not name_list[0].startswith("-"):
-                    self.layers[command.PROG][1].append(completer)
-                else:
-                    for name in name_list:
-                        self.layers[command.PROG][2][name] = completer
-                        option_names.append(name)
-            self.layers[command.PROG][0] = WordCompleter(
-                option_names + ["--help", "-h"]
-            )
+            if command.ARGS is not None:
+                for name_list, descr in command.ARGS.items():
+                    name_list = name_list.split(",")
+                    if descr[0] == Complete.CHOICES:
+                        completer = WordCompleter(descr[3]["choices"])
+                    elif descr[0] == Complete.LOCAL_FILE:
+                        completer = local_file_completer
+                    elif descr[0] == Complete.REMOTE_FILE:
+                        completer = remote_file_completer
+                    elif descr[0] == Complete.NONE:
+                        completer = None
+                    if len(name_list) == 1 and not name_list[0].startswith("-"):
+                        self.layers[command.PROG][1].append(completer)
+                    else:
+                        for name in name_list:
+                            self.layers[command.PROG][2][name] = completer
+                            option_names.append(name)
+                self.layers[command.PROG][0] = WordCompleter(
+                    option_names + ["--help", "-h"]
+                )
 
         self.completer = WordCompleter(list(self.layers))
 

pwncat/commands/base.py

@@ -122,7 +122,10 @@ class CommandDefinition:
     PROG = "unimplemented"
     """ The name of your new command """
     ARGS = {}
-    """ A dictionary of parameter definitions created with the ``parameter`` function. """
+    """ A dictionary of parameter definitions created with the ``parameter`` function.
+    If this is None, your command will receive the raw argument string and no processing
+    will be done except removing the leading command name.
+    """
     DEFAULTS = {}
     """ A dictionary of default values (passed directly to ``ArgumentParser.set_defaults``) """
     LOCAL = False
@@ -146,9 +149,13 @@ class CommandDefinition:
         into an argparse object. """
 
         # Create the parser object
-        self.parser = argparse.ArgumentParser(prog=self.PROG, description=self.__doc__)
-
-        self.build_parser(self.parser, self.ARGS)
+        if self.ARGS is not None:
+            self.parser = argparse.ArgumentParser(
+                prog=self.PROG, description=self.__doc__
+            )
+            self.build_parser(self.parser, self.ARGS)
+        else:
+            self.parser = None
 
     def run(self, args):
         """

pwncat/commands/help.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+import textwrap
+
 import pwncat
 from pwncat.commands import CommandParser
 from pwncat.commands.base import CommandDefinition, Complete, parameter
@@ -16,7 +18,10 @@ class Command(CommandDefinition):
         if args.topic:
             for command in pwncat.victim.command_parser.commands:
                 if command.PROG == args.topic:
-                    command.parser.print_help()
+                    if command.parser is not None:
+                        command.parser.print_help()
+                    else:
+                        print(textwrap.dedent(command.__doc__).strip())
                     break
         else:
             util.info("the following commands are available:")

pwncat/commands/local.py

@@ -6,14 +6,11 @@ from pwncat.commands.base import parameter
 
 
 class Command(CommandDefinition):
+    """ Run a local shell command on your attacking machine """
 
     PROG = "local"
-    ARGS = {
-        "argv": parameter(
-            Complete.NONE, nargs="+", help="the local shell command to run"
-        )
-    }
+    ARGS = None
     LOCAL = True
 
     def run(self, args):
-        subprocess.run(args.argv, shell=True)
+        subprocess.run(args, shell=True)

pwncat/commands/run.py

@@ -6,13 +6,16 @@ from pwncat.commands.base import CommandDefinition, Complete, parameter
 
 
 class Command(CommandDefinition):
+    """
+    Run a shell command on the victim host and display the output.
+    
+    **NOTE** This must be a non-interactive command. If an interactive command
+        is run, you will have to use C-c to return to the pwncat prompt and then
+        C-d to get back to your interactive remote prompt in order to interact
+        with the remote host again!"""
 
     PROG = "run"
-    ARGS = {
-        "argv": parameter(
-            Complete.NONE, nargs="+", help="The command to run on the remote host"
-        )
-    }
+    ARGS = None
 
     def run(self, args):
-        sys.stdout.buffer.write(pwncat.victim.run(args.argv))
+        sys.stdout.buffer.write(pwncat.victim.run(args))