Port to aglib: entry (#2567)

* start: Port to aglib

* start: fix message on missing entry

pwndbg/commands/__init__.py

@@ -9,6 +9,7 @@ from typing import Any
 from typing import Callable
 from typing import Dict
 from typing import List
+from typing import Literal
 from typing import Optional
 from typing import Set
 from typing import Tuple
@@ -344,6 +345,32 @@ def OnlyWithArch(arch_names: List[str]) -> Callable[[Callable[P, T]], Callable[P
     return decorator
 
 
+def OnlyWithDbg(
+    *dbg_names: Literal["lldb", "gdb"],
+) -> Callable[[Callable[P, T]], Callable[P, Optional[T]]]:
+    """Decorates function to work only with the specified debugger."""
+
+    def decorator(function: Callable[P, T]) -> Callable[P, Optional[T]]:
+        @functools.wraps(function)
+        def _OnlyWithDbg(*a: P.args, **kw: P.kwargs) -> Optional[T]:
+            if pwndbg.dbg.is_gdblib_available():
+                if "gdb" in dbg_names:
+                    return function(*a, **kw)
+            else:
+                if "lldb" in dbg_names:
+                    return function(*a, **kw)
+
+            dbg_str = ", ".join(dbg_names)
+            log.error(
+                f"{function.__name__}: This command may only be run on the {dbg_str} debugger(s)"
+            )
+            return None
+
+        return _OnlyWithDbg
+
+    return decorator
+
+
 def OnlyWithKernelDebugSyms(function: Callable[P, T]) -> Callable[P, Optional[T]]:
     @functools.wraps(function)
     def _OnlyWithKernelDebugSyms(*a: P.args, **kw: P.kwargs) -> Optional[T]:
@@ -713,7 +740,6 @@ def load_commands() -> None:
         import pwndbg.commands.segments
         import pwndbg.commands.shell
         import pwndbg.commands.slab
-        import pwndbg.commands.start
         import pwndbg.commands.tips
         import pwndbg.commands.version
 
@@ -761,6 +787,7 @@ def load_commands() -> None:
     import pwndbg.commands.search
     import pwndbg.commands.sigreturn
     import pwndbg.commands.spray
+    import pwndbg.commands.start
     import pwndbg.commands.telescope
     import pwndbg.commands.tls
     import pwndbg.commands.valist

pwndbg/commands/start.py

@@ -9,25 +9,33 @@ import argparse
 from argparse import RawTextHelpFormatter
 from shlex import quote
 
-import gdb
-
 import pwndbg
+import pwndbg.aglib.elf
+import pwndbg.aglib.proc
+import pwndbg.color.message as M
 import pwndbg.commands
-import pwndbg.gdblib.elf
-import pwndbg.gdblib.symbol
+import pwndbg.dbg
 from pwndbg.commands import CommandCategory
-from pwndbg.dbg import EventType
+from pwndbg.dbg import BreakpointLocation
+
+if pwndbg.dbg.is_gdblib_available():
+    import gdb
+
+
+def breakpoint_at_entry():
+    addr = int(pwndbg.aglib.elf.entry())
+    if not addr:
+        print(M.error("No entry address found for the binary."))
+        return
 
-break_on_first_instruction = False
+    proc = pwndbg.dbg.selected_inferior()
+    bp = proc.break_at(BreakpointLocation(addr), internal=True)
 
+    async def ctrl(ec: pwndbg.dbg_mod.ExecutionController):
+        await ec.cont(bp)
+        bp.remove()
 
-@pwndbg.dbg.event_handler(EventType.START)
-def on_start() -> None:
-    global break_on_first_instruction
-    if break_on_first_instruction:
-        spec = "*%#x" % (int(pwndbg.gdblib.elf.entry()))
-        gdb.Breakpoint(spec, temporary=True)
-        break_on_first_instruction = False
+    proc.dispatch_execution_controller(ctrl)
 
 
 # Starting from 3rd paragraph, the description is
@@ -58,6 +66,7 @@ parser.add_argument(
 
 
 @pwndbg.commands.ArgparsedCommand(parser, aliases=["main", "init"], category=CommandCategory.START)
+@pwndbg.commands.OnlyWithDbg("gdb")
 def start(args=None) -> None:
     if args is None:
         args = []
@@ -66,8 +75,7 @@ def start(args=None) -> None:
     symbols = ["main", "_main", "start", "_start", "init", "_init"]
 
     for symbol in symbols:
-        address = pwndbg.gdblib.symbol.address(symbol)
-
+        address = pwndbg.dbg.selected_inferior().symbol_address_from_name(symbol)
         if not address:
             continue
 
@@ -103,23 +111,37 @@ To start the inferior without using a shell, use "set startup-with-shell off".
 """,
 )
 parser.add_argument(
-    "args", nargs="*", type=str, default=[], help="The arguments to run the binary with."
+    "args", nargs="*", type=str, default=None, help="The arguments to run the binary with."
 )
 
 
 @pwndbg.commands.ArgparsedCommand(parser, category=CommandCategory.START)
 @pwndbg.commands.OnlyWithFile
-def entry(args=[]) -> None:
-    global break_on_first_instruction
-    break_on_first_instruction = True
-    run = "run " + " ".join(map(quote, args))
-    gdb.execute(run, from_tty=False)
+def entry(args=None) -> None:
+    if args is None:
+        args = []
+
+    if pwndbg.dbg.is_gdblib_available():
+        run = "starti " + " ".join(map(quote, args))
+        gdb.execute(run, from_tty=False)
+    else:
+        # TODO: LLDB, In the future, we should handle `run -s` here to automate setup.
+        # For now, we only support stopping at the entry breakpoint.
+        if not pwndbg.aglib.proc.alive:
+            print(
+                M.error(
+                    "The program is not running. Start the program with `run -s` and then use `entry` to set the breakpoint."
+                )
+            )
+            return
+    breakpoint_at_entry()
 
 
 @pwndbg.commands.ArgparsedCommand(
     "Alias for 'tbreak __libc_start_main; run'.", category=CommandCategory.START
 )
 @pwndbg.commands.OnlyWithFile
+@pwndbg.commands.OnlyWithDbg("gdb")
 def sstart() -> None:
     gdb.Breakpoint("__libc_start_main", temporary=True)
     gdb.execute("run")