diff --git a/CHANGELOG.md b/CHANGELOG.md index a51a771e..0828e137 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 The Changelog starts with v0.4.1, because we did not keep one before that, and simply didn't have the time to go back and retroactively create one. +## [Unreleased] + +### Changed +- Added `force` argument to `Manager.load_modules` to enable reloading modules. +- Added the `Manager.reload_module` method to reload a specific module. +- Added `--force/-f` and `--reload/-r` arguments to the `load` command ([#241](https://github.com/calebstewart/pwncat/issues/231)). +- Added the `--reload/-r` argument to the `run` command to reload modules prior to executing. + ## [0.5.4] - 2022-01-27 Bug fix for the `load` command. diff --git a/pwncat/commands/load.py b/pwncat/commands/load.py index 2fbe082b..f0dbeae3 100644 --- a/pwncat/commands/load.py +++ b/pwncat/commands/load.py @@ -1,5 +1,7 @@ #!/usr/bin/env python3 +from pathlib import Path + import pwncat from pwncat.commands import Complete, Parameter, CommandDefinition @@ -21,11 +23,30 @@ class Command(CommandDefinition): Complete.LOCAL_FILE, help="Path to a python package directory to load modules from", nargs="+", - ) + ), + "--force,-f": Parameter( + Complete.NONE, + help="Force loading the given module(s) even if they were already loaded.", + action="store_true", + default=False, + ), + "--reload,-r": Parameter( + Complete.NONE, + help="Synonym for --force", + action="store_true", + dest="force", + ), } DEFAULTS = {} LOCAL = True def run(self, manager: "pwncat.manager.Manager", args): - manager.load_modules(*args.path) + # Python's pkgutil.walk_packages doesn't produce an error + # if the path doesn't exist, so we double check that each + # provided path exists prior to calling it. + for path in args.path: + if not Path(path).expanduser().exists(): + self.parser.error(f"{path}: no such file or directory") + + manager.load_modules(*args.path, force=args.force) diff --git a/pwncat/commands/run.py b/pwncat/commands/run.py index 0dd95615..1ea74d26 100644 --- a/pwncat/commands/run.py +++ b/pwncat/commands/run.py @@ -20,13 +20,22 @@ class Command(CommandDefinition): To locate available modules, you can use the `search` command. To find documentation on individual modules including expected arguments, you can use the `info` command. + + The `--reload/-r` argument can be used to force the module to be + reloaded from disk prior to execution. This is useful for debugging + mostly. """ PROG = "run" ARGS = { - "--raw,-r": Parameter( + "--raw": Parameter( Complete.NONE, action="store_true", help="Display raw results unformatted" ), + "--reload,-r": Parameter( + Complete.NONE, + action="store_true", + help="Reload the module prior to execution", + ), "--traceback,-t": Parameter( Complete.NONE, action="store_true", help="Show traceback for module errors" ), @@ -50,6 +59,18 @@ def run(self, manager: "pwncat.manager.Manager", args): elif args.module is None: module_name = manager.config.module.name + if args.reload: + if args.module is not None: + # We have a target-specific module name, so we need to + # use find_module to locate the specific module. If this + # module doesn't exist, this does nothing, and we will error + # out below. + for module in manager.target.find_module(args.module, exact=True): + manager.reload_module(module) + else: + # Reload the module from the config context + manager.reload_module(manager.config.module) + # Parse key=value pairs values = {} for arg in args.args: @@ -68,12 +89,6 @@ def run(self, manager: "pwncat.manager.Manager", args): if args.module is not None: manager.config.back() - except pwncat.modules.ModuleFailed as exc: - if args.traceback: - console.print_exception() - else: - console.log(f"[red]error[/red]: module failed: {exc}") - return except pwncat.modules.ModuleNotFound: console.log(f"[red]error[/red]: {module_name}: not found") return @@ -86,6 +101,12 @@ def run(self, manager: "pwncat.manager.Manager", args): except pwncat.modules.InvalidArgument as exc: console.log(f"[red]error[/red]: invalid argument: {exc}") return + except pwncat.modules.ModuleFailed as exc: + if args.traceback: + console.print_exception() + else: + console.log(f"[red]error[/red]: module failed: {exc}") + return if isinstance(result, list): result = [r for r in result if not r.hidden] diff --git a/pwncat/manager.py b/pwncat/manager.py index 393cb012..b5e9d0e6 100644 --- a/pwncat/manager.py +++ b/pwncat/manager.py @@ -921,32 +921,131 @@ def create_db_session(self): return self.db.open() - def load_modules(self, *paths): + def load_modules(self, *paths, force: bool = False): """Dynamically load modules from the specified paths If a module has the same name as an already loaded module, it will take it's place in the module list. This includes built-in modules. + + If the fully qualified Python module name matches a built-in or + previously loaded module, then you will need to provide the `force` + argument. This forces a reload of the given python modules. + + :param force: force reloading the module if already loaded + :type force: bool """ for loader, module_name, _ in pkgutil.walk_packages( - paths, prefix="pwncat.modules." + [os.path.expanduser(p) for p in paths], prefix="pwncat.modules." ): - - # Why is this check *not* part of pkgutil??????? D:< - if module_name not in sys.modules: - module = loader.find_module(module_name).load_module(module_name) - else: + # Strip off the prefix + name = module_name.split("pwncat.modules.")[1] + + # If the module was already loaded and the user specified force, + # delete the reference from `sys.modules` so that the module is + # reloaded. + if module_name in sys.modules and force: + del sys.modules[module_name] + + # This should happen automatically through pkgutils... + # but it didn't seem to be, so I just double check it here. + if module_name in sys.modules: module = sys.modules[module_name] + else: + module = loader.find_module(module_name).load_module(module_name) if getattr(module, "Module", None) is None: continue - # Create an instance of this module - module_name = module_name.split("pwncat.modules.")[1] - self.modules[module_name] = module.Module() + # Create a module instance + new_module = module.Module() + setattr(new_module, "name", name) + setattr(new_module, "_loader", loader) + setattr(new_module, "_module_name", module_name) + + # Save the old module + old_module = self.modules.get(name, new_module) + + # Update the module list + self.modules[name] = new_module - # Store it's name so we know it later - setattr(self.modules[module_name], "name", module_name) + # Update the current module context if we just replaced a module + if old_module != new_module and self.config.module == old_module: + # Save the local configuration for the module + config = self.config.locals + # Use the new module + self.config.use(new_module) + + # Attempt to re-set any configuration items previously set + # This could fail if the argument definitions changed on disk. + for key, value in config.items(): + try: + self.config.set(key, value) + except ValueError as exc: + self.log( + f"[yellow]warning[/yellow]: failed to re-set module config: {key}: {exc}" + ) + + def reload_module( + self, module: Union[str, "pwncat.modules.BaseModule"] + ) -> "pwncat.modules.BaseModule": + """Reload the given module from disk. The module can either be an + instance of a module returned from :py:meth:`~Session.find_module` + or a fully-qualified module name. If the selected module is currently + being used in a module context, the context will be switched to the + reloaded module, and all parameters will be reset (potentially causing + errors if the module arguments have changed). + + :param module: fully-qualified module name or module object + :type module: Union[str, pwncat.modules.BaseModule] + """ + + # Locate the module to reload if passed as a string + if isinstance(module, str): + module = self.modules[module] + + # "module" is ambiguous since Python uses it, so we save + # the module objects in old_module and new_module. + old_module = module + + # Saved loader and python module name from the old module. + name = old_module.name + loader = old_module._loader + module_name = old_module._module_name + + # Remove the python module from sys.modules so the loader + # will re-load it from disk. + if module_name in sys.modules: + del sys.modules[module_name] + + # Load the python module + module = loader.find_module(module_name).load_module(module_name) + + # Construct a pwncat module instance + new_module = module.Module() + setattr(new_module, "name", name) + setattr(new_module, "_loader", loader) + setattr(new_module, "_module_name", module_name) + + # Store the new module + self.modules[name] = new_module + + # Replace the current context if it was set to the old module + if self.config.module == old_module: + # Save the local configuration for the module + config = self.config.locals + # Use the new module + self.config.use(new_module) + + # Attempt to re-set any configuration items previously set + # This could fail if the argument definitions changed on disk. + for key, value in config.items(): + try: + self.config.set(key, value) + except ValueError as exc: + self.log( + f"[yellow]warning[/yellow]: failed to re-set module config: {key}: {exc}" + ) def log(self, *args, **kwargs): """Output a log entry"""