diff --git a/pwncat/commands/__init__.py b/pwncat/commands/__init__.py index ca17a0d..a4f1aef 100644 --- a/pwncat/commands/__init__.py +++ b/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)) diff --git a/pwncat/commands/base.py b/pwncat/commands/base.py index ae06f05..b84e5fd 100644 --- a/pwncat/commands/base.py +++ b/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): """ diff --git a/pwncat/commands/help.py b/pwncat/commands/help.py index de78b85..329c878 100644 --- a/pwncat/commands/help.py +++ b/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:") diff --git a/pwncat/commands/local.py b/pwncat/commands/local.py index 30f8a9c..30e8385 100644 --- a/pwncat/commands/local.py +++ b/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) diff --git a/pwncat/commands/run.py b/pwncat/commands/run.py index 0db4b2e..c7c1771 100644 --- a/pwncat/commands/run.py +++ b/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))