Home directory prefix support

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`
 

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

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:

pwncat/commands/download.py

@@ -14,16 +14,18 @@
 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):
     """Download a file from the remote host to the local host"""
 
     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):

pwncat/commands/upload.py

@@ -13,18 +13,17 @@
 
 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):
     """Upload a file from the local host to the remote host"""
 
     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
         ),
     }