Added module reloading API and CLI interface

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.
 

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)

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]

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"""