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
+            if command.parser:
                 args = command.parser.parse_args(args)
+            else:
+                args = line
 
             # Run the command
             command.run(args)
@@ -315,6 +320,7 @@ class CommandLexer(RegexLexer):
         for command in commands:
             root.append(("^" + re.escape(command.PROG), Name.Function, command.PROG))
             mode = []
+            if command.ARGS is not None:
                 for args, descr in command.ARGS.items():
                     for arg in args.split(","):
                         if not arg.startswith("-"):
@@ -408,6 +414,7 @@ class CommandCompleter(Completer):
         for command in commands:
             self.layers[command.PROG] = [None, [], {}]
             option_names = []
+            if command.ARGS is not None:
                 for name_list, descr in command.ARGS.items():
                     name_list = name_list.split(",")
                     if descr[0] == Complete.CHOICES:

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__)
-
+        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:
+                    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))