diff --git a/CHANGELOG.md b/CHANGELOG.md index 2d18d09b..d1fbde86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,9 @@ and simply didn't have the time to go back and retroactively create one. - Added licensing for pwncat (MIT) - Added background listener API and commands ([#43](https://github.com/calebstewart/pwncat/issues/43)) - Added Windows privilege escalation via BadPotato plugin ([#106](https://github.com/calebstewart/pwncat/issues/106)) +- Added command parameter parsers +- Added homedir (`~`) support on local file completer and parser +- Added homedir (`~`) support on remote file completer and parser ### Removed - Removed `setup.py` and `requirements.txt` diff --git a/pwncat/commands/__init__.py b/pwncat/commands/__init__.py index 3d7dc2bd..0b0693f6 100644 --- a/pwncat/commands/__init__.py +++ b/pwncat/commands/__init__.py @@ -97,6 +97,44 @@ class Complete(Enum): """ Do not provide argument completions """ +class ParseType(Enum): + """ + Command type. This defines how command parameter arguments are parsed + """ + + NONE = auto() + """ No specific type given, so no interpreter needed """ + LOCAL_FILE = auto() + """ Local file type """ + REMOTE_FILE = auto() + """ Remote file type """ + + +class ParameterParse: + def parse(value, session: "pwncat.manager.Session"): + """This is where the parameter will be parsed""" + raise NotImplementedError + + +class LocalFileParse(ParameterParse): + def parse(value, session: "pwncat.manager.Session"): + if value.startswith("~"): + return os.path.expanduser(value) + return value + + +class RemoteFileParse(ParameterParse): + def parse(value, session: "pwncat.manager.Session"): + if value.startswith("~/") or value.startswith("~\\"): + homedir = session.platform.getenv("HOME") + if not homedir: + """Windows support""" + homedir = session.platform.getenv("USERPROFILE") + if homedir: + return homedir + value[1:] + return value + + class StoreConstOnce(argparse.Action): """Only allow the user to store a value in the destination once. This prevents users from selection multiple actions in the privesc parser.""" @@ -180,6 +218,8 @@ class Parameter: :param complete: the completion type :type complete: Complete + :param parser: the parsing type + :type parser: ParseType :param token: the Pygments token to highlight this argument with :type token: Pygments Token :param group: true for a group definition, a string naming the group to be a part of, or none @@ -191,12 +231,14 @@ class Parameter: def __init__( self, complete: Complete, + parser=ParseType.NONE, token=token.Name.Label, group: str = None, *args, **kwargs, ): self.complete = complete + self.parser = parser self.token = token self.group = group self.args = args @@ -338,6 +380,27 @@ def __iter__(wself): parser.set_defaults(**self.DEFAULTS) + def parse_args(self, args, fallback): + if not self.parser: + return fallback + + parsed = vars(self.parser.parse_args(args)) + for [argkey, argobj] in self.ARGS.items(): + if argkey not in parsed or argobj.parser is ParseType.NONE: + continue + + if argobj.parser is not ParseType.NONE: + parser = None + if argobj.parser is ParseType.LOCAL_FILE: + parser = LocalFileParse + elif argobj.parser is ParseType.REMOTE_FILE: + parser = RemoteFileParse + + if parser is not None: + parsed[argkey] = parser.parse(parsed[argkey], self.manager.target) + + return argparse.Namespace(**parsed) + def resolve_blocks(source: str): """This is a dumb lexer that turns strings of text with code blocks (squigly @@ -659,10 +722,7 @@ def dispatch_line(self, line: str, prog_name: str = None): prog_name = temp_name # Parse the arguments - if command.parser: - args = command.parser.parse_args(args) - else: - args = line + args = command.parse_args(args, line) # Run the command command.run(self.manager, args) @@ -853,6 +913,14 @@ def get_completions(self, document: Document, complete_event: CompleteEvent): if path == "": path = "." + if path.startswith("~"): + homedir = self.manager.target.platform.getenv("HOME") + if not homedir: + """Windows support""" + homedir = self.manager.target.platform.getenv("USERPROFILE") + if homedir: + path = homedir + path[1:] + for name in self.manager.target.platform.listdir(path): if name.startswith(partial_name): yield Completion( @@ -873,6 +941,9 @@ def get_completions(self, document: Document, complete_event: CompleteEvent): if path == "": path = "." + if path.startswith("~"): + path = os.path.expanduser(path) + # Ensure the directory exists if not os.path.isdir(path): return diff --git a/pwncat/commands/connect.py b/pwncat/commands/connect.py index 11f2d7ae..07a52045 100644 --- a/pwncat/commands/connect.py +++ b/pwncat/commands/connect.py @@ -203,8 +203,6 @@ def run(self, manager: "pwncat.manager.Manager", args): console.log("[red]error[/red]: multiple ports specified") return - console.log(args.pos_port) - if args.port is not None: query_args["port"] = args.port if args.pos_port is not None: diff --git a/pwncat/commands/download.py b/pwncat/commands/download.py index 002636ce..16e6ad39 100644 --- a/pwncat/commands/download.py +++ b/pwncat/commands/download.py @@ -14,7 +14,7 @@ import pwncat from pwncat import util from pwncat.util import console -from pwncat.commands import Complete, Parameter, CommandDefinition +from pwncat.commands import Complete, Parameter, ParseType, CommandDefinition class Command(CommandDefinition): @@ -22,8 +22,10 @@ class Command(CommandDefinition): PROG = "download" ARGS = { - "source": Parameter(Complete.REMOTE_FILE), - "destination": Parameter(Complete.LOCAL_FILE, nargs="?"), + "source": Parameter(Complete.REMOTE_FILE, parser=ParseType.REMOTE_FILE), + "destination": Parameter( + Complete.LOCAL_FILE, nargs="?", parser=ParseType.LOCAL_FILE + ), } def run(self, manager: "pwncat.manager.Manager", args): diff --git a/pwncat/commands/upload.py b/pwncat/commands/upload.py index 13e547d1..e6da9c68 100644 --- a/pwncat/commands/upload.py +++ b/pwncat/commands/upload.py @@ -13,7 +13,7 @@ import pwncat from pwncat.util import console, copyfileobj, human_readable_size, human_readable_delta -from pwncat.commands import Complete, Parameter, CommandDefinition +from pwncat.commands import Complete, Parameter, ParseType, CommandDefinition class Command(CommandDefinition): @@ -21,10 +21,9 @@ class Command(CommandDefinition): PROG = "upload" ARGS = { - "source": Parameter(Complete.LOCAL_FILE), + "source": Parameter(Complete.LOCAL_FILE, parser=ParseType.LOCAL_FILE), "destination": Parameter( - Complete.REMOTE_FILE, - nargs="?", + Complete.REMOTE_FILE, nargs="?", parser=ParseType.REMOTE_FILE ), } diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b