Add --shell argument to pwndbg command (#1233)

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__
 

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):
+def pwndbg_(filter_pattern, shell, all_):
-    for name, docs in list_and_filter_commands(filter_pattern):
+    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__
 

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:

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."
+    "loaded %i pwndbg commands and %i shell commands. Type %s for a list."
-    % (len(pwndbg.commands.commands), message.notice("pwndbg [filter]")),
+    % (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,
 )
 

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)",
 ]
 

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

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