Fix reload command

pwndbg/__init__.py

@@ -31,40 +31,6 @@ import pwndbg.ui
 __version__ = pwndbg.lib.version.__version__
 version = __version__
 
-__all__ = [
-    "arch",
-    "auxv",
-    "chain",
-    "color",
-    "disasm",
-    "dt",
-    "elf",
-    "enhance",
-    "events",
-    "file",
-    "function",
-    "heap",
-    "hexdump",
-    "ida",
-    "info",
-    "leakfind",
-    "linkmap",
-    "malloc",
-    "memoize",
-    "memory",
-    "p2p",
-    "proc",
-    "regs",
-    "remote",
-    "search",
-    "stack",
-    "strings",
-    "symbol",
-    "typeinfo",
-    "ui",
-    "vmmap",
-]
-
 from pwndbg.gdblib import prompt
 
 prompt.set_prompt()

pwndbg/commands/__init__.py

@@ -64,6 +64,11 @@ def list_current_commands():
 
 GDB_BUILTIN_COMMANDS = list_current_commands()
 
+# Set in `reload` command so that we can skip double checking for registration
+# of an already existing command when re-registering GDB CLI commands
+# (there is no way to unregister a command in GDB 12.x)
+pwndbg_is_reloading = getattr(gdb, "pwndbg_is_reloading", False)
+
 
 class Command(gdb.Command):
     """Generic command wrapper"""
@@ -99,6 +104,7 @@ class Command(gdb.Command):
         if (
             command_name in GDB_BUILTIN_COMMANDS
             and command_name not in self.builtin_override_whitelist
+            and not pwndbg_is_reloading
         ):
             raise Exception('Cannot override non-whitelisted built-in command "%s"' % command_name)
 

pwndbg/commands/reload.py

@@ -1,5 +1,7 @@
-import types
-from imp import reload as _reload
+import importlib
+import sys
+
+import gdb
 
 import pwndbg
 import pwndbg.commands
@@ -8,23 +10,19 @@ import pwndbg.lib.memoize
 from pwndbg.commands import CommandCategory
 
 
-def rreload(module, mdict=None) -> None:
-    """Recursively reload modules."""
-    name = module.__name__
-
-    if mdict is None:
-        mdict = []
+def rreload(module, _exclude_mods=None) -> None:
+    """Recursively reload modules.
+    Impl based on https://stackoverflow.com/a/66661311/1508881"""
+    delete_folders = ["pwndbg"]
 
-    for attribute_name in getattr(module, "__all__", []) or []:
-        attribute = getattr(module, attribute_name, None)
-        if isinstance(attribute, types.ModuleType) and attribute not in mdict:
-            mdict.append(attribute)
-            rreload(attribute, mdict)
+    for module in list(sys.modules.keys()):
+        if any(folder in module for folder in delete_folders):
+            del sys.modules[module]
 
-    try:
-        _reload(module)
-    except Exception as e:
-        pass
+    # Mark that we are reloading; this is used to prevent ArgparsedCommand from
+    # erroring out on re-registering the same commands we had registered before
+    gdb.pwndbg_is_reloading = True
+    importlib.import_module("pwndbg")
 
 
 @pwndbg.commands.ArgparsedCommand("Reload pwndbg.", category=CommandCategory.PWNDBG)