diff --git a/pwndbg/commands/__init__.py b/pwndbg/commands/__init__.py index 227619cf0..4e7a2b1b1 100644 --- a/pwndbg/commands/__init__.py +++ b/pwndbg/commands/__init__.py @@ -54,7 +54,9 @@ class Command(gdb.Command): builtin_override_whitelist = {"up", "down", "search", "pwd", "start"} history = {} # type: Dict[int,str] - def __init__(self, function, prefix=False, command_name=None): + def __init__(self, function, prefix=False, command_name=None, shell=False): + self.shell = shell + if command_name is None: command_name = function.__name__ diff --git a/pwndbg/commands/misc.py b/pwndbg/commands/misc.py index 2e09b926f..725543682 100644 --- a/pwndbg/commands/misc.py +++ b/pwndbg/commands/misc.py @@ -64,6 +64,11 @@ parser = argparse.ArgumentParser( Prints out a list of all pwndbg commands. The list can be optionally filtered if filter_pattern is passed. """ ) + +group = parser.add_mutually_exclusive_group() +group.add_argument("--shell", action="store_true", help="Only display shell commands") +group.add_argument("--all", dest="all_", action="store_true", help="Only display shell commands") + parser.add_argument( "filter_pattern", type=str, @@ -74,8 +79,18 @@ parser.add_argument( @pwndbg.commands.ArgparsedCommand(parser, command_name="pwndbg") -def pwndbg_(filter_pattern): - for name, docs in list_and_filter_commands(filter_pattern): +def pwndbg_(filter_pattern, shell, all_): + if all_: + shell_cmds = True + pwndbg_cmds = True + elif shell: + shell_cmds = True + pwndbg_cmds = False + else: + shell_cmds = False + pwndbg_cmds = True + + for name, docs in list_and_filter_commands(filter_pattern, pwndbg_cmds, shell_cmds): print("%-20s %s" % (name, docs)) @@ -98,7 +113,7 @@ def distance(a, b): ) -def list_and_filter_commands(filter_str): +def list_and_filter_commands(filter_str, pwndbg_cmds=True, shell_cmds=False): sorted_commands = list(pwndbg.commands.commands) sorted_commands.sort(key=lambda x: x.__name__) @@ -108,6 +123,14 @@ def list_and_filter_commands(filter_str): results = [] for c in sorted_commands: + # If this is a shell command and we didn't ask for shell commands, skip it + if c.shell and not shell_cmds: + continue + + # If this is a normal command and we didn't ask for normal commands, skip it + if not c.shell and not pwndbg_cmds: + continue + name = c.__name__ docs = c.__doc__ diff --git a/pwndbg/commands/shell.py b/pwndbg/commands/shell.py index 3eb9cd785..e63b6093f 100644 --- a/pwndbg/commands/shell.py +++ b/pwndbg/commands/shell.py @@ -81,7 +81,7 @@ def register_shell_function(cmd, deprecated=False): handler.__name__ = str(cmd) handler.__doc__ = doc - pwndbg.commands.Command(handler, False) + pwndbg.commands.Command(handler, shell=True) for cmd in pwncmds: diff --git a/pwndbg/gdblib/prompt.py b/pwndbg/gdblib/prompt.py index 9f5caab19..cd87f419c 100644 --- a/pwndbg/gdblib/prompt.py +++ b/pwndbg/gdblib/prompt.py @@ -14,9 +14,11 @@ funcs_list_str = ", ".join( message.notice("$" + f.name) for f in pwndbg.gdbutils.functions.functions ) +num_pwndbg_cmds = sum(1 for _ in filter(lambda c: not c.shell, pwndbg.commands.commands)) +num_shell_cmds = sum(1 for _ in filter(lambda c: c.shell, pwndbg.commands.commands)) hint_lines = ( - "loaded %i commands. Type %s for a list." - % (len(pwndbg.commands.commands), message.notice("pwndbg [filter]")), + "loaded %i pwndbg commands and %i shell commands. Type %s for a list." + % (num_pwndbg_cmds, num_shell_cmds, message.notice("pwndbg [--shell | --all] [filter]")), "created %s gdb functions (can be used with print/break)" % funcs_list_str, ) diff --git a/tests/test_loads.py b/tests/test_loads.py index bf1832266..669b35d90 100644 --- a/tests/test_loads.py +++ b/tests/test_loads.py @@ -7,7 +7,7 @@ from .utils import compile_binary from .utils import run_gdb_with_script HELLO = [ - "pwndbg: loaded ### commands. Type pwndbg [filter] for a list.", + "pwndbg: loaded ### pwndbg commands and ### shell commands. Type pwndbg [--shell | --all] [filter] for a list.", "pwndbg: created $rebase, $ida gdb functions (can be used with print/break)", ] diff --git a/tests/test_misc.py b/tests/test_misc.py index a2b6734e0..069238748 100644 --- a/tests/test_misc.py +++ b/tests/test_misc.py @@ -17,7 +17,7 @@ def test_list_and_filter_commands_filter(): def test_list_and_filter_commands_full_list(): - all_commands = list_and_filter_commands("") + all_commands = list_and_filter_commands("", pwndbg_cmds=True, shell_cmds=True) def get_doc(c): return c.__doc__.strip().splitlines()[0] if c.__doc__ else None @@ -26,3 +26,27 @@ def test_list_and_filter_commands_full_list(): cmd_name_docs.sort() assert all_commands == cmd_name_docs + + +def test_list_and_filter_commands_shell(): + all_commands = list_and_filter_commands("", pwndbg_cmds=False, shell_cmds=True) + + def get_doc(c): + return c.__doc__.strip().splitlines()[0] if c.__doc__ else None + + cmd_name_docs = [(c.__name__, get_doc(c)) for c in pwndbg.commands.commands if c.shell] + cmd_name_docs.sort() + + assert all_commands == cmd_name_docs + + +def test_list_and_filter_commands_pwndbg_cmds(): + all_commands = list_and_filter_commands("", pwndbg_cmds=True, shell_cmds=False) + + def get_doc(c): + return c.__doc__.strip().splitlines()[0] if c.__doc__ else None + + cmd_name_docs = [(c.__name__, get_doc(c)) for c in pwndbg.commands.commands if not c.shell] + cmd_name_docs.sort() + + assert all_commands == cmd_name_docs diff --git a/tests/utils.py b/tests/utils.py index a5d753fdc..85d08f909 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -40,7 +40,11 @@ def run_gdb_with_script(binary="", core="", pybefore=None, pyafter=None, timeout # The pwndbg banner shows number of loaded commands, it might differ between # testing environments, so lets change it to ### - output = re.sub(r"loaded [0-9]+ commands", r"loaded ### commands", output) + output = re.sub( + r"loaded [0-9]+ pwndbg commands and [0-9]+ shell commands", + r"loaded ### pwndbg commands and ### shell commands", + output, + ) return output