diff --git a/pwndbg/commands/__init__.py b/pwndbg/commands/__init__.py index ff61f4d47..fd2bc1020 100644 --- a/pwndbg/commands/__init__.py +++ b/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, ) diff --git a/pwndbg/commands/heap.py b/pwndbg/commands/heap.py index f9101935a..53ae1d268 100644 --- a/pwndbg/commands/heap.py +++ b/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) diff --git a/pwndbg/commands/misc.py b/pwndbg/commands/misc.py index dc45fceda..216711680 100644 --- a/pwndbg/commands/misc.py +++ b/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 diff --git a/tests/gdb-tests/tests/test_misc.py b/tests/gdb-tests/tests/test_misc.py index d96099c7f..36f7e75c6 100644 --- a/tests/gdb-tests/tests/test_misc.py +++ b/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