Home directory prefix support

CHANGELOG.md

@@ -21,6 +21,8 @@ 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
 ### Removed
 - Removed `setup.py` and `requirements.txt`
 

pwncat/commands/__init__.py

@@ -46,6 +46,7 @@ def run(self, manager: "pwncat.manager.Manager", args: "argparse.Namespace"):
 from io import TextIOWrapper
 from enum import Enum, auto
 from typing import Dict, List, Type, Callable, Iterable
+from pathlib import Path
 from functools import partial
 
 import rich.text
@@ -97,6 +98,31 @@ 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()
+ """ File type """
+
+
+class ParameterParse:
+ def parse(value):
+ """This is where the parameter will be parsed"""
+ raise NotImplementedError
+
+
+class LocalFileParse:
+ def parse(value):
+ if value.startswith("~"):
+ homedir = str(Path.home())
+ 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 +206,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 +219,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 +368,18 @@ 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 ParseType.LOCAL_FILE:
+ parsed[argkey] = LocalFileParse.parse(parsed[argkey])
+ 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 +701,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)
@@ -873,6 +912,10 @@ def get_completions(self, document: Document, complete_event: CompleteEvent):
  if path == "":
  path = "."
 
+ if path.startswith("~"):
+ homedir = str(Path.home())
+ path = homedir + path[1:]
+
  # Ensure the directory exists
  if not os.path.isdir(path):
  return

pwncat/commands/upload.py

@@ -13,15 +13,15 @@
 
 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="?",