Add support for command categories (#1477)

pwndbg/commands/__init__.py

@@ -53,10 +53,18 @@ class Command(gdb.Command):
     history = {}  # type: Dict[int,str]
 
     def __init__(
-        self, function, prefix=False, command_name=None, shell=False, is_alias=False, aliases=[]
+        self,
+        function,
+        prefix=False,
+        command_name=None,
+        shell=False,
+        is_alias=False,
+        aliases=[],
+        category="Misc",
     ) -> None:
         self.is_alias: bool = is_alias
         self.aliases = aliases
+        self.category = category
         self.shell: bool = shell
 
         if command_name is None:
@@ -422,7 +430,15 @@ def OnlyWithResolvedHeapSyms(function):
 
 class _ArgparsedCommand(Command):
     def __init__(
-        self, parser, function, command_name=None, is_alias=False, aliases=[], *a, **kw
+        self,
+        parser,
+        function,
+        command_name=None,
+        is_alias=False,
+        aliases=[],
+        category="Misc",
+        *a,
+        **kw,
     ) -> None:
         self.parser = parser
         if command_name is None:
@@ -438,7 +454,13 @@ class _ArgparsedCommand(Command):
         function.__doc__ = self.parser.description.strip()
 
         super(_ArgparsedCommand, self).__init__(
-            function, command_name=command_name, is_alias=is_alias, aliases=aliases, *a, **kw
+            function,
+            command_name=command_name,
+            is_alias=is_alias,
+            aliases=aliases,
+            category=category,
+            *a,
+            **kw,
         )
 
     def split_args(self, argument):
@@ -449,7 +471,7 @@ class _ArgparsedCommand(Command):
 class ArgparsedCommand:
     """Adds documentation and offloads parsing for a Command via argparse"""
 
-    def __init__(self, parser_or_desc, aliases=[], command_name=None) -> None:
+    def __init__(self, parser_or_desc, aliases=[], command_name=None, category="Misc") -> None:
         """
         :param parser_or_desc: `argparse.ArgumentParser` instance or `str`
         """
@@ -459,7 +481,7 @@ class ArgparsedCommand:
             self.parser = parser_or_desc
         self.aliases = aliases
         self._command_name = command_name
-
+        self.category = category
         # We want to run all integer and otherwise-unspecified arguments
         # through fix() so that GDB parses it.
         for action in self.parser._actions:
@@ -472,9 +494,15 @@ class ArgparsedCommand:
 
     def __call__(self, function):
         for alias in self.aliases:
-            _ArgparsedCommand(self.parser, function, command_name=alias, is_alias=True)
+            _ArgparsedCommand(
+                self.parser, function, command_name=alias, is_alias=True, category=self.category
+            )
         return _ArgparsedCommand(
-            self.parser, function, command_name=self._command_name, aliases=self.aliases
+            self.parser,
+            function,
+            command_name=self._command_name,
+            aliases=self.aliases,
+            category=self.category,
         )
 
 

pwndbg/commands/heap.py

@@ -123,7 +123,7 @@ parser.add_argument(
 )
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -155,7 +155,7 @@ Default to the current thread's arena.""",
 parser.add_argument("addr", nargs="?", type=int, default=None, help="Address of the arena.")
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -174,7 +174,7 @@ def arena(addr=None) -> None:
 parser = argparse.ArgumentParser(description="List this process's arenas.")
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -194,7 +194,7 @@ Default to the current thread's tcache.""",
 parser.add_argument("addr", nargs="?", type=int, default=None, help="Address of the tcache.")
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -211,7 +211,7 @@ def tcache(addr=None) -> None:
 parser = argparse.ArgumentParser(description="Print the mp_ struct's contents.")
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -230,7 +230,7 @@ Default to current thread's arena.""",
 parser.add_argument("addr", nargs="?", type=int, default=None, help="Address of the arena.")
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -261,7 +261,7 @@ parser.add_argument(
 )
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -346,7 +346,7 @@ parser.add_argument("addr", nargs="?", type=int, default=None, help="Address of
 parser.add_argument("tcache_addr", nargs="?", type=int, default=None, help="Address of the tcache.")
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -372,7 +372,7 @@ parser.add_argument("addr", nargs="?", type=int, default=None, help="Address of
 parser.add_argument("verbose", nargs="?", type=bool, default=True, help="Show extra detail.")
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -403,7 +403,7 @@ parser.add_argument("addr", nargs="?", type=int, default=None, help="Address of
 parser.add_argument("verbose", nargs="?", type=bool, default=True, help="Show extra detail.")
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -434,7 +434,7 @@ parser.add_argument("addr", nargs="?", type=int, default=None, help="Address of
 parser.add_argument("verbose", nargs="?", type=bool, default=False, help="Show extra detail.")
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -465,7 +465,7 @@ parser.add_argument("addr", nargs="?", type=int, default=None, help="Address of
 parser.add_argument("verbose", nargs="?", type=bool, default=False, help="Show extra detail.")
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -500,7 +500,7 @@ parser.add_argument(
 )
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -537,7 +537,7 @@ parser.add_argument(
 )
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -655,7 +655,7 @@ parser.add_argument(
 )
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
@@ -831,7 +831,7 @@ try_free_parser = argparse.ArgumentParser(
 try_free_parser.add_argument("addr", nargs="?", help="Address passed to free")
 
 
-@pwndbg.commands.ArgparsedCommand(try_free_parser)
+@pwndbg.commands.ArgparsedCommand(try_free_parser, category="Heap")
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWhenHeapIsInitialized
 def try_free(addr) -> None:
@@ -1208,7 +1208,7 @@ parser.add_argument(
 )
 
 
-@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.ArgparsedCommand(parser, category="Heap")
 def heap_config(filter_pattern) -> None:
     display_config(filter_pattern, "heap", has_file_command=False)
 

pwndbg/commands/misc.py

@@ -1,5 +1,6 @@
 import argparse
 import errno
+from collections import defaultdict
 
 import gdb
 
@@ -87,8 +88,10 @@ def pwndbg_(filter_pattern, shell, all_) -> None:
 
     from tabulate import tabulate
 
-    table_data = []
-    for name, aliases, docs in list_and_filter_commands(filter_pattern, pwndbg_cmds, shell_cmds):
+    table_data = defaultdict(lambda: [])
+    for name, aliases, category, docs in list_and_filter_commands(
+        filter_pattern, pwndbg_cmds, shell_cmds
+    ):
         alias_str = ""
         aliases_len = 0
         if aliases:
@@ -96,11 +99,16 @@ def pwndbg_(filter_pattern, shell, all_) -> None:
             alias_str = f" [{', '.join(aliases)}]"
 
         command_names = C.green(name) + alias_str
-        table_data.append((command_names, docs))
+        table_data[category].append((command_names, docs))
 
-    print(
-        tabulate(table_data, headers=[f"{C.green('Command')} [{C.blue('Aliases')}]", "Description"])
-    )
+    for category, data in table_data.items():
+        category_header = category + " Commands"
+        print(
+            tabulate(
+                data, headers=[f"{C.green(category_header)} [{C.blue('Aliases')}]", "Description"]
+            )
+        )
+        print()
 
 
 parser = argparse.ArgumentParser(description="Print the distance between the two arguments.")
@@ -153,6 +161,6 @@ def list_and_filter_commands(filter_str, pwndbg_cmds=True, shell_cmds=False):
             docs = docs.splitlines()[0]
 
         if not filter_str or filter_str in name.lower() or (docs and filter_str in docs.lower()):
-            results.append((name, c.aliases, docs))
+            results.append((name, c.aliases, c.category, docs))
 
     return results

tests/gdb-tests/tests/test_misc.py

@@ -4,12 +4,12 @@ import pwndbg.commands
 from pwndbg.commands.misc import list_and_filter_commands
 
 STACK_COMMANDS = [
-    ("canary", [], "Print out the current stack canary."),
-    ("context", ["ctx"], "Print out the current register, instruction, and stack context."),
-    ("down", [], "Select and print stack frame called by this one."),
-    ("retaddr", [], "Print out the stack addresses that contain return addresses."),
-    ("stack", [], "Dereferences on stack data with specified count and offset."),
-    ("up", [], "Select and print stack frame that called this one."),
+    ("canary", [], "Misc", "Print out the current stack canary."),
+    ("context", ["ctx"], "Misc", "Print out the current register, instruction, and stack context."),
+    ("down", [], "Misc", "Select and print stack frame called by this one."),
+    ("retaddr", [], "Misc", "Print out the stack addresses that contain return addresses."),
+    ("stack", [], "Misc", "Dereferences on stack data with specified count and offset."),
+    ("up", [], "Misc", "Select and print stack frame that called this one."),
 ]
 
 
@@ -31,7 +31,7 @@ def test_list_and_filter_commands_full_list(pwndbg_cmds, shell_cmds):
     if shell_cmds:
         commands.extend([c for c in pwndbg.commands.commands if not c.is_alias and c.shell])
 
-    cmd_name_docs = [(c.__name__, c.aliases, get_doc(c)) for c in commands]
+    cmd_name_docs = [(c.__name__, c.aliases, c.category, get_doc(c)) for c in commands]
     cmd_name_docs.sort()
 
     assert all_commands == cmd_name_docs