Improve pwndbg-lldb REPL (#2625)

* add pwndbg-lldb history file

* add pwndbg-lldb last_command

* fix arrow up key, history command

* fix arrow up key, history command

pwndbg/color/message.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Callable
+
 import pwndbg.lib.config
 from pwndbg import config
 from pwndbg.color import generateColorFunction
@@ -105,3 +107,12 @@ def prompt(msg: object) -> str:
 
 def alive_prompt(msg: object) -> str:
     return generateColorFunction(config.prompt_alive_color)(msg)
+
+
+def readline_escape(func_message: Callable[[str], str], text: str) -> str:
+    # For readline-based applications, non-printable escape codes must be
+    # wrapped with special markers (\001 and \002). These markers inform
+    # readline to ignore the escape sequences when calculating the prompt's width.
+    # Without these markers, the prompt may break when navigating command history
+    # with the UP arrow key or for long commands.
+    return "\x01" + func_message("\x02" + text + "\x01") + "\x02"

pwndbg/dbg/lldb/repl/__init__.py

@@ -59,8 +59,8 @@ from pwndbg.dbg.lldb.repl.io import get_io_driver
 from pwndbg.dbg.lldb.repl.proc import EventHandler
 from pwndbg.dbg.lldb.repl.proc import ProcessDriver
 from pwndbg.dbg.lldb.repl.readline import PROMPT
-from pwndbg.dbg.lldb.repl.readline import disable_readline
 from pwndbg.dbg.lldb.repl.readline import enable_readline
+from pwndbg.dbg.lldb.repl.readline import wrap_with_history
 from pwndbg.lib.tips import color_tip
 from pwndbg.lib.tips import get_tip_of_the_day
 
@@ -168,6 +168,7 @@ def show_greeting() -> None:
         print(colored_tip)
 
 
+@wrap_with_history
 def run(startup: List[str] | None = None, debug: bool = False) -> None:
     """
     Runs the Pwndbg REPL under LLDB. Optionally enters the commands given in
@@ -200,6 +201,8 @@ def run(startup: List[str] | None = None, debug: bool = False) -> None:
     signal.signal(signal.SIGINT, handle_sigint)
 
     show_greeting()
+    last_command = ""
+
     while True:
         # Execute the prompt hook and ask for input.
         dbg._fire_prompt_hook()
@@ -211,6 +214,11 @@ def run(startup: List[str] | None = None, debug: bool = False) -> None:
                 startup_i += 1
             else:
                 line = input(PROMPT)
+                # If the input is empty (i.e., 'Enter'), use the previous command
+                if line:
+                    last_command = line
+                else:
+                    line = last_command
         except EOFError:
             # Exit the REPL if there's nothing else to run.
             print()
@@ -402,29 +410,6 @@ def run(startup: List[str] | None = None, debug: bool = False) -> None:
             break
 
 
-def make_pty() -> Tuple[str, int]:
-    """
-    We need to make a pseudo-terminal ourselves if we want the process to handle
-    naturally for the user. Returns a tuple with the filaname and the file
-    descriptor if successful.
-    """
-    import ctypes
-
-    libc = ctypes.CDLL("libc.so.6")
-    pty = libc.posix_openpt(2)
-    if pty <= 0:
-        return None
-
-    libc.ptsname.restype = ctypes.c_char_p
-    name = libc.ptsname(pty)
-
-    if libc.unlockpt(pty) != 0:
-        libc.close(pty)
-        return None
-
-    return name, pty
-
-
 def parse(args: List[str], parser: argparse.ArgumentParser, unsupported: List[str]) -> Any | None:
     """
     Parses a list of string arguments into an object containing the parsed

pwndbg/dbg/lldb/repl/readline.py

@@ -6,13 +6,25 @@ Mostly concerns itself with argument completion.
 
 from __future__ import annotations
 
+import contextlib
+import functools
+import os.path
+from typing import Callable
+from typing import ParamSpec
+from typing import TypeVar
+
 import gnureadline as readline
 import lldb
 
 from pwndbg.color import message
 from pwndbg.dbg.lldb import LLDB
 
-PROMPT = message.prompt("pwndbg-lldb> ")
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+PROMPT = message.readline_escape(message.prompt, "pwndbg-lldb> ")
+HISTORY_FILE = os.path.expanduser("~/.pwndbg_history")
 
 complete_values = lldb.SBStringList()
 complete_descrs = lldb.SBStringList()
@@ -63,6 +75,27 @@ def display_completions(substitutions, matches, longest_match_len):
     print(readline.get_line_buffer(), end="", flush=True)
 
 
+def wrap_with_history(function: Callable[P, T]) -> Callable[P, T]:
+    @functools.wraps(function)
+    def _wrapped(*a: P.args, **kw: P.kwargs) -> T:
+        with ctx_with_history():
+            return function(*a, **kw)
+
+    return _wrapped
+
+
+@contextlib.contextmanager
+def ctx_with_history():
+    readline.set_history_length(1000)
+    if os.path.exists(HISTORY_FILE):
+        readline.read_history_file(HISTORY_FILE)
+
+    try:
+        yield
+    finally:
+        readline.write_history_file(HISTORY_FILE)
+
+
 def enable_readline(dbg: LLDB):
     """
     Enables the readline functionality.

pwndbg/gdblib/prompt.py

@@ -117,12 +117,10 @@ def set_prompt() -> None:
     prompt = "pwndbg> "
 
     if not disable_colors:
-        prompt = "\x02" + prompt + "\x01"  # STX + prompt + SOH
         if pwndbg.aglib.proc.alive:
-            prompt = message.alive_prompt(prompt)
+            prompt = message.readline_escape(message.alive_prompt, prompt)
         else:
-            prompt = message.prompt(prompt)
-        prompt = "\x01" + prompt + "\x02"  # SOH + prompt + STX
+            prompt = message.readline_escape(message.prompt, prompt)
 
     gdb.execute(f"set prompt {prompt}")