Improve argument parsing and launch handling in LLDB Pwndbg (#3081)

* Improve argument parsing in LLDB Pwndbg

* Use `driver.has_connection()` to determine remote status

pwndbg-lldb.py

@@ -2,12 +2,30 @@
 
 from __future__ import annotations
 
+import argparse
 import os
 import re
 import subprocess
 import sys
 from typing import List
 
+PARSER = argparse.ArgumentParser(prog="pwndbg-lldb")
+PARSER.add_argument("-v", "--verbose", action="store_true", help="Enable debug output")
+PARSER.add_argument("target", nargs="?")
+parser_attach = PARSER.add_mutually_exclusive_group()
+parser_attach.add_argument(
+    "-n", "--attach-name", help="Tells the debugger to attach to a process with the given name."
+)
+parser_attach.add_argument(
+    "-p", "--attach-pid", help="Tells the debugger to attach to a process with the given pid."
+)
+PARSER.add_argument(
+    "-w",
+    "--wait-for",
+    action="store_true",
+    help="Tells the debugger to wait for a process with the given pid or name to launch before attaching.",
+)
+
 
 def find_lldb_version() -> List[int]:
     """
@@ -41,7 +59,8 @@ def find_lldb_python_path() -> str:
 
 
 if __name__ == "__main__":
-    debug = "PWNDBG_LLDB_DEBUG" in os.environ
+    args = PARSER.parse_args()
+    debug = args.verbose
 
     # Find the path for the LLDB Python bindings.
     path = find_lldb_python_path()
@@ -96,21 +115,50 @@ if __name__ == "__main__":
         print("[-] Launcher: Initializing Pwndbg")
     lldbinit.main(debugger, lldb_version[0], lldb_version[1], debug=debug)
 
-    # Run our REPL until the user decides to leave.
-    if len(sys.argv) > 2:
-        print(f"Usage: {sys.argv[0]} [filename]", file=sys.stderr)
-        sys.exit(1)
-
-    target = None
-    if len(sys.argv) == 2:
-        target = sys.argv[1]
-
     from pwndbg.dbg.lldb.repl import PwndbgController
+    from pwndbg.dbg.lldb.repl import print_error
+    from pwndbg.dbg.lldb.repl import print_warn
     from pwndbg.dbg.lldb.repl import run as run_repl
 
     if debug:
         print("[-] Launcher: Entering Pwndbg CLI")
 
+    # Prepare the startup commands.
+    startup = []
+    if args.target:
+        # DEVIATION: The LLDB CLI silently ignores any target information passed
+        # to it when using either '--attach-name' or '--attach-pid', but Pwndbg
+        # unconditionally uses it, with a warning.
+        startup = [f"target create '{args.target}'"]
+
+    if args.attach_name is not None:
+        wait = "--waitfor" if args.wait_for else ""
+        startup.append(f'process attach --name "{args.attach_name}" {wait}')
+    elif args.attach_pid is not None:
+        # DEVIATION: While the LLDB CLI accepts '--wait-for' in combination with
+        # both '--attach-name' and '--attach-pid', it silently ignores it when
+        # used with the latter. Pwndbg prints out a warning, instead.
+        if args.wait_for:
+            print_warn("'--wait-for' has no effect when used with '--attach-pid'")
+
+        startup.append(f'process attach --pid "{args.attach_pid}"')
+    else:
+        if args.wait_for:
+            # Ideally, we would have `ArgumentParser` do this for us, but
+            # nesting argument groups has been deprecated since Python 3.11, and
+            # the deprecation message suggests it was never even supported in
+            # the first place :/
+            print_error(
+                "'--wait-for' must be used in combination with either '--attach-name' or '--attach-pid'"
+            )
+            PARSER.print_usage()
+            sys.exit(1)
+
+    if (args.attach_pid is not None or args.attach_name is not None) and args.target:
+        print_warn(
+            "have both a target and an attach request, your target may be overwritten on attach"
+        )
+
     def drive(startup: List[str] | None):
         async def drive(c: PwndbgController):
             if startup is not None:
@@ -122,7 +170,7 @@ if __name__ == "__main__":
 
         return drive
 
-    run_repl(drive([f"target create '{target}'"] if target else None), debug=debug)
+    run_repl(drive(startup), debug=debug)
 
     # Dispose of our debugger and terminate LLDB.
     lldb.SBDebugger.Destroy(debugger)

pwndbg/dbg/lldb/repl/__init__.py

@@ -385,8 +385,9 @@ def exec_repl_command(
 
     # Let the user get an LLDB prompt if they so desire.
     if bits[0] == "lldb":
+        print_warn("You are now entering LLDB mode. To exit, type 'quit', 'exit' or Ctrl-D.")
         print_warn(
-            "You are now entering LLDB mode. In this mode, certain commands may cause Pwndbg to break. Proceed with caution."
+            "In this mode, certain commands may cause Pwndbg to break. Proceed with caution."
         )
         dbg.debugger.RunCommandInterpreter(
             True, False, lldb.SBCommandInterpreterRunOptions(), 0, False, False
@@ -493,8 +494,8 @@ def exec_repl_command(
             pass
 
     if bits[0].startswith("r") and "run".startswith(bits[0]):
-        # `run` is an alias for `process launch`
+        # `run` is an alias for `process launch -X true --`
-        process_launch(driver, relay, bits[1:], dbg)
+        process_launch(driver, relay, ["-X", "true", "--"] + bits[1:], dbg)
         return True
 
     if bits[0] == "c" or (bits[0].startswith("con") and "continue".startswith(bits[0])):
@@ -652,11 +653,48 @@ def exec_repl_command(
     return True
 
 
-def parse(args: List[str], parser: argparse.ArgumentParser, unsupported: List[str]) -> Any | None:
+def _bool_of_string(val: str) -> bool:
+    """
+    Convert a string to a boolean value.
+
+    For use with ArgumentParser.
+    """
+    if val.lower() in ("true", "1", "yes"):
+        return True
+    elif val.lower() in ("false", "0", "no"):
+        return False
+    else:
+        raise ValueError(f"{val} is not a recognized boolean value")
+
+
+def parse(
+    args: List[str],
+    parser: argparse.ArgumentParser,
+    unsupported: List[str],
+    raw_marker: str | None = None,
+) -> Any | None:
     """
     Parses a list of string arguments into an object containing the parsed
     data.
+
+    If `raw_marker` is not `None`, the argument list will be split in
+    two, with all arguments before the split being fed to the argument parser,
+    and all arguments after the split being returned as-is. In this case the
+    return value is a tuple.
     """
+    raw = None
+    if raw_marker is not None:
+        # Always return something, even if we match nothing.
+        raw = []
+
+        try:
+            index = args.index(raw_marker)
+            raw = args[index + 1 :]
+            args = args[:index]
+        except ValueError:
+            # Ugly, but avoids going over the list twice.
+            pass
+
     try:
         args = parser.parse_args(args)
     except SystemExit:
@@ -673,9 +711,54 @@ def parse(args: List[str], parser: argparse.ArgumentParser, unsupported: List[st
             print_error(f"Pwndbg does not support --{unsup} yet")
             return None
 
+    if raw is not None:
+        # If called with `raw_marked`, return a tuple.
+        return args, raw
+
     return args
 
 
+class AutoTarget:
+    """
+    During the execution of some commands, the LLDB CLI automatically creates an
+    empty target and selects it before the command is executed.
+    """
+
+    def __init__(self, dbg: LLDB):
+        self.error = lldb.SBError()
+        self._dbg = dbg
+        self._created_target = False
+
+        count = dbg.debugger.GetNumTargets()
+        if count == 0:
+            # Create the target.
+            self.target = dbg.debugger.CreateTarget(None, None, None, True, self.error)
+
+            if not self.error.success:
+                return
+
+            # On success, select it and remember that it has been created.
+            dbg.debugger.SetSelectedTarget(self.target)
+            self._created_target = True
+        elif count == 1:
+            # Just use the current target.
+            self.target = dbg.debugger.GetTargetAtIndex(0)
+            assert self.target, f"SBDebugger::GetNumTargets() is 1, but SBDebugger::GetTargetAtIndex(0) is {self.target}"
+        else:
+            raise AssertionError(
+                f"Pwndbg does not support multiple targets, so SBDebugger::GetNumTargets() must always be 0 or 1, but is {count}"
+            )
+
+    def __bool__(self):
+        return self.error.success
+
+    def close(self):
+        if self._created_target:
+            assert self._dbg.debugger.DeleteTarget(
+                self.target
+            ), "Could not delete the target we've just created. What?"
+
+
 def run_ipython_shell():
     @contextmanager
     def switch_to_ipython_env():
@@ -707,7 +790,7 @@ def run_ipython_shell():
         start_ipi()
 
 
-target_create_ap = argparse.ArgumentParser(add_help=False)
+target_create_ap = argparse.ArgumentParser(add_help=False, prog="target create")
 target_create_ap.add_argument("-S", "--sysroot")
 target_create_ap.add_argument("-a", "--arch")
 target_create_ap.add_argument("-b", "--build")
@@ -796,12 +879,12 @@ def target_create(args: List[str], dbg: LLDB) -> None:
     return
 
 
-process_launch_ap = argparse.ArgumentParser(add_help=False)
+process_launch_ap = argparse.ArgumentParser(add_help=False, prog="process launch")
 process_launch_ap.add_argument("-A", "--disable-aslr")
 process_launch_ap.add_argument("-C", "--script-class")
-process_launch_ap.add_argument("-E", "--environment")
+process_launch_ap.add_argument("-E", "--environment", action="append")
 process_launch_ap.add_argument("-P", "--plugin")
-process_launch_ap.add_argument("-X", "--shell-expand-args")
+process_launch_ap.add_argument("-X", "--shell-expand-args", type=_bool_of_string)
 process_launch_ap.add_argument("-a", "--arch")
 process_launch_ap.add_argument("-c", "--shell")
 process_launch_ap.add_argument("-e", "--stderr")
@@ -817,9 +900,7 @@ process_launch_ap.add_argument("run-args", nargs="*")
 process_launch_unsupported = [
     "disable-aslr",
     "script-class",
-    "environment",
     "plugin",
-    "shell-expand-args",
     "arch",
     "shell",
     "stderr",
@@ -837,9 +918,15 @@ def process_launch(driver: ProcessDriver, relay: EventRelay, args: List[str], db
     """
     Launches a process with the given arguments.
     """
-    args = parse(args, process_launch_ap, process_launch_unsupported)
+    result = parse(args, process_launch_ap, process_launch_unsupported, raw_marker="--")
-    if not args:
+    if result is None:
         return
+    args, raw = result
+
+    launch_args = getattr(args, "run-args", []) + raw
+    if args.shell_expand_args:
+        # Perform shell expansion.
+        launch_args = [os.path.expanduser(os.path.expandvars(arg)) for arg in launch_args]
 
     targets = dbg.debugger.GetNumTargets()
     assert targets < 2
@@ -852,8 +939,9 @@ def process_launch(driver: ProcessDriver, relay: EventRelay, args: List[str], db
         return
 
     target: lldb.SBTarget = dbg.debugger.GetTargetAtIndex(0)
-    # Make sure the LLDB driver knows that this is a local process.
+
-    dbg._current_process_is_gdb_remote = False
+    # Make sure LLDB knows the correct remote or local status of this launch.
+    dbg._current_process_is_gdb_remote = driver.has_connection()
 
     if target.GetPlatform().GetName() == "qemu-user":
         # Force qemu-user as remote, pwndbg depends on that, eg: for download procfs files
@@ -863,8 +951,9 @@ def process_launch(driver: ProcessDriver, relay: EventRelay, args: List[str], db
     result = driver.launch(
         target,
         io_driver,
-        [f"{name}={value}" for name, value in os.environ.items()],
+        [f"{name}={value}" for name, value in os.environ.items()]
-        getattr(args, "run-args", []),
+        + (args.environment if args.environment else []),
+        launch_args,
         os.getcwd(),
     )
 
@@ -896,7 +985,7 @@ def process_launch(driver: ProcessDriver, relay: EventRelay, args: List[str], db
         dbg._trigger_event(EventType.STOP)
 
 
-process_attach_ap = argparse.ArgumentParser(add_help=False)
+process_attach_ap = argparse.ArgumentParser(add_help=False, prog="process attach")
 process_attach_ap.add_argument("-C", "--python-class")
 process_attach_ap.add_argument("-P", "--plugin")
 process_attach_ap.add_argument("-c", "--continue", action="store_true")
@@ -920,11 +1009,7 @@ def _attach_with_info(
     """
     Attaches to a process based on SBAttachInfo information
     """
-    targets = dbg.debugger.GetNumTargets()
+    assert dbg.debugger.GetNumTargets() < 2
-    assert targets < 2
-    if targets == 0:
-        print_error("no target, create one using the 'target create' command")
-        return
 
     # TODO/FIXME: This should ask:
     # 'There is a running process, detach from it and attach?: [Y/n]'
@@ -934,14 +1019,24 @@ def _attach_with_info(
 
     io_driver = get_io_driver()
 
+    auto = AutoTarget(dbg)
+    if not auto:
+        print_error(f"could not create empty target for attaching: {auto.error.description}")
+        auto.close()
+        return
+
+    # Make sure LLDB knows the correct remote or local status of this attach.
+    dbg._current_process_is_gdb_remote = driver.has_connection()
+
     result = driver.attach(
-        dbg.debugger.GetTargetAtIndex(0),
+        auto.target,
         io_driver,
         info,
     )
 
     if not result.success:
         print_error(f"could not attach to process: {result.description}")
+        auto.close()
         return
 
     # Continue execution if the user has requested it.
@@ -1006,7 +1101,7 @@ def attach(driver: ProcessDriver, relay: EventRelay, args: List[str], dbg: LLDB)
     _attach_with_info(driver, relay, dbg, info)
 
 
-process_connect_ap = argparse.ArgumentParser(add_help=False)
+process_connect_ap = argparse.ArgumentParser(add_help=False, prog="process connect")
 process_connect_ap.add_argument("-p", "--plugin")
 process_connect_ap.add_argument("remoteurl")
 
@@ -1029,53 +1124,21 @@ def process_connect(driver: ProcessDriver, relay: EventRelay, args: List[str], d
         print_error("debugger is already connected")
         return
 
-    target = dbg.debugger.GetSelectedTarget()
-    error = lldb.SBError()
-    created_target = False
-    if target is None or not target.IsValid():
-        # Create a new, empty target, the same way the LLDB command line would.
-        #
-        # The LLDB command line sets the default triple based on the
-        # architecture value set in the `target.default-arch` setting. We do the
-        # same.
-        try:
-            result = dbg._execute_lldb_command("settings show target.default-arch")
-
-            # The result of this command has the following form:
-            #
-            # (lldb) settings show target.default-arch
-            # target.default-arch (arch) = <value>
-            #
-            # Where <value> may be empty, for no value.
-            arch = result.split("=")[1].strip()
-        except pwndbg.dbg_mod.Error:
-            arch = ""
-
-        triple = f"{arch}-unknown-unknown" if len(arch) > 0 else None
-
-        target = dbg.debugger.CreateTarget(None, triple, None, True, error)
-        if not error.success or not target.IsValid():
-            print_error(
-                f"could not automatically create target for 'process connect': {error.description}"
-            )
-            return
-
-        dbg.debugger.SetSelectedTarget(target)
-        created_target = True
-
     # Make sure the LLDB driver knows that this is a remote process.
     dbg._current_process_is_gdb_remote = True
 
+    auto = AutoTarget(dbg)
+    if not auto:
+        print_error(f"could not create empty target for connection: {auto.error.description}")
+        auto.close()
+        return
+
     io_driver = get_io_driver()
-    error = driver.connect(target, io_driver, args.remoteurl, "gdb-remote")
+    error = driver.connect(auto.target, io_driver, args.remoteurl, "gdb-remote")
 
     if not error.success:
         print_error(f"could not connect to remote process: {error.description}")
-        if created_target:
+        auto.close()
-            # Delete the target we previously created.
-            assert dbg.debugger.DeleteTarget(
-                target
-            ), "Could not delete the target we've just created. What?"
         return
 
     # Tell the debugger that the process was suspended, if there is a process.
@@ -1083,7 +1146,7 @@ def process_connect(driver: ProcessDriver, relay: EventRelay, args: List[str], d
         dbg._trigger_event(EventType.STOP)
 
 
-gdb_remote_ap = argparse.ArgumentParser(add_help=False)
+gdb_remote_ap = argparse.ArgumentParser(add_help=False, prog="gdb-remote")
 gdb_remote_ap.add_argument("remoteurl")
 
 
@@ -1120,7 +1183,7 @@ def gdb_remote(driver: ProcessDriver, relay: EventRelay, args: List[str], dbg: L
     process_connect(driver, relay, ["-p", "gdb-remote", f"connect://{url}:{port}"], dbg)
 
 
-continue_ap = argparse.ArgumentParser(add_help=False)
+continue_ap = argparse.ArgumentParser(add_help=False, prog="continue")
 continue_ap.add_argument("-i", "--ignore-count")
 continue_unsupported = ["ignore-count"]