Add initialization under LLDB (#2253)

Dockerfile.lldb

@@ -0,0 +1,67 @@
+# This dockerfile was created for development & testing purposes, for APT-based distro.
+#
+# Build as:             docker build -t pwndbg .
+#
+# For testing use:      docker run --rm -it --cap-add=SYS_PTRACE --security-opt seccomp=unconfined pwndbg bash
+#
+# For development, mount the directory so the host changes are reflected into container:
+#   docker run -it --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v `pwd`:/pwndbg pwndbg bash
+#
+
+ARG image=mcr.microsoft.com/devcontainers/base:jammy
+FROM $image
+
+WORKDIR /pwndbg
+
+ENV PIP_NO_CACHE_DIR=true
+ENV LANG en_US.utf8
+ENV TZ=America/New_York
+ENV ZIGPATH=/opt/zig
+ENV PWNDBG_VENV_PATH=/venv
+
+RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && \
+    echo $TZ > /etc/timezone && \
+    apt-get update && \
+    apt-get install -y locales && \
+    rm -rf /var/lib/apt/lists/* && \
+    localedef -i en_US -c -f UTF-8 -A /usr/share/locale/locale.alias en_US.UTF-8 && \
+    apt-get update && \
+    apt-get install -y vim
+
+ADD ./setup.sh /pwndbg/
+ADD ./poetry.lock /pwndbg/
+ADD ./pyproject.toml /pwndbg/
+ADD ./poetry.toml /pwndbg/
+
+# pyproject.toml requires these files, pip install would fail
+RUN touch README.md && mkdir pwndbg && touch pwndbg/empty.py
+
+RUN DEBIAN_FRONTEND=noninteractive ./setup.sh
+
+# Cleanup dummy files
+RUN rm README.md && rm -rf pwndbg
+
+# Comment these lines if you won't run the tests.
+ADD ./setup-dev.sh /pwndbg/
+RUN ./setup-dev.sh
+
+ADD . /pwndbg/
+
+ARG LOW_PRIVILEGE_USER="vscode"
+
+# Add .gdbinit to the home folder of both root and vscode users (if vscode user exists)
+# This is useful for a VSCode dev container, not really for test builds
+RUN if [ ! -f ~/.gdbinit ]; then echo "source /pwndbg/gdbinit.py" >> ~/.gdbinit; fi && \
+    if id -u ${LOW_PRIVILEGE_USER} > /dev/null 2>&1; then \
+        su ${LOW_PRIVILEGE_USER} -c 'if [ ! -f ~/.gdbinit ]; then echo "source /pwndbg/gdbinit.py" >> ~/.gdbinit; fi'; \
+    fi
+
+RUN apt-get install -y lldb-16
+
+# Add .lldbinit to the home folder of both root and vscode users (if vscode user exists)
+# This is useful for a VSCode dev container, not really for test builds
+RUN if [ ! -f ~/.lldbinit ]; then echo "command script import /pwndbg/lldbinit.py" >> ~/.lldbinit; fi && \
+    if id -u ${LOW_PRIVILEGE_USER} > /dev/null 2>&1; then \
+        su ${LOW_PRIVILEGE_USER} -c 'if [ ! -f ~/.lldbinit ]; then echo "command script import /pwndbg/lldbinit.py" >> ~/.lldbinit; fi'; \
+    fi
+

docker-compose.yml

@@ -36,6 +36,16 @@ services:
       args:
         image: debian:11
 
+  lldb:
+    <<: *base-spec
+    build:
+      context: .
+      dockerfile: Dockerfile.lldb
+      args:
+        image: debian:12
+    volumes:
+      - .:/pwndbg
+
   archlinux:
     <<: *base-spec
     build:

lint.sh

@@ -84,5 +84,5 @@ vermin -vvv --no-tips -t=3.8- --eval-annotations --violations ${LINT_FILES}
 
 # mypy is run in a separate step on GitHub Actions
 if [[ -z "$GITHUB_ACTIONS" ]]; then
-    mypy pwndbg gdbinit.py
+    mypy pwndbg gdbinit.py lldbinit.py
 fi

lldbinit.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+import cProfile
+import hashlib
+import os
+import shutil
+import site
+import subprocess
+import sys
+import time
+from glob import glob
+from pathlib import Path
+from typing import List
+from typing import Tuple
+
+import lldb
+
+
+def hash_file(file_path: str | Path) -> str:
+    with open(file_path, "rb") as f:
+        file_hash = hashlib.sha256()
+        while True:
+            chunk = f.read(8192)
+            if not chunk:
+                break
+            file_hash.update(chunk)
+    return file_hash.hexdigest()
+
+
+def run_poetry_install(poetry_path: os.PathLike[str], dev: bool = False) -> Tuple[str, str, int]:
+    command: List[str | os.PathLike[str]] = [poetry_path, "install"]
+    if dev:
+        command.extend(("--with", "dev"))
+    result = subprocess.run(command, capture_output=True, text=True)
+    return result.stdout.strip(), result.stderr.strip(), result.returncode
+
+
+def find_poetry() -> Path | None:
+    poetry_path = shutil.which("poetry")
+    if poetry_path is not None:
+        return Path(poetry_path)
+
+    # On some systems `poetry` is installed in "~/.local/bin/" but this directory is
+    # not on the $PATH
+    poetry_path = Path("~/.local/bin/poetry").expanduser()
+    if poetry_path.exists():
+        return poetry_path
+
+    return None
+
+
+def is_dev_mode(venv_path: Path) -> bool:
+    # If "dev.marker" exists in the venv directory, the user ran setup-dev.sh and is
+    # considered a developer
+    return (venv_path / "dev.marker").exists()
+
+
+def update_deps(src_root: Path, venv_path: Path) -> None:
+    poetry_lock_hash_path = venv_path / "poetry.lock.hash"
+
+    current_hash = hash_file(src_root / "poetry.lock")
+    stored_hash = None
+    if poetry_lock_hash_path.exists():
+        stored_hash = poetry_lock_hash_path.read_text().strip()
+
+    # If the hashes don't match, update the dependencies
+    if current_hash != stored_hash:
+        poetry_path = find_poetry()
+        if poetry_path is None:
+            print(
+                "Poetry was not found on the $PATH. Please ensure it is installed and on the path, "
+                "or run `./setup.sh` to manually update Python dependencies."
+            )
+            return
+
+        dev_mode = is_dev_mode(venv_path)
+        stdout, stderr, return_code = run_poetry_install(poetry_path, dev=dev_mode)
+        if return_code == 0:
+            poetry_lock_hash_path.write_text(current_hash)
+
+            # Only print the poetry output if anything was actually updated
+            if "No dependencies to install or update" not in stdout:
+                print(stdout)
+        else:
+            print(stderr, file=sys.stderr)
+
+
+def fixup_paths(src_root: Path, venv_path: Path):
+    site_pkgs_path = glob(str(venv_path / "lib/*/site-packages"))[0]
+
+    # add virtualenv's site-packages to sys.path and run .pth files
+    site.addsitedir(site_pkgs_path)
+
+    # remove existing, system-level site-packages from sys.path
+    for site_packages in site.getsitepackages():
+        if site_packages in sys.path:
+            sys.path.remove(site_packages)
+
+    # Set virtualenv's bin path (needed for utility tools like ropper, pwntools etc)
+    bin_path = str(venv_path / "bin")
+    os.environ["PATH"] = bin_path + os.pathsep + os.environ.get("PATH", "")
+
+    # Add pwndbg directory to sys.path so it can be imported
+    sys.path.insert(0, str(src_root))
+
+    # Push virtualenv's site-packages to the front
+    if site_pkgs_path in sys.path:
+        sys.path.remove(site_pkgs_path)
+    sys.path.insert(1, site_pkgs_path)
+
+
+def get_venv_path(src_root: Path):
+    venv_path_env = os.environ.get("PWNDBG_VENV_PATH")
+    if venv_path_env:
+        return Path(venv_path_env).expanduser().resolve()
+    else:
+        return src_root / ".venv"
+
+
+def skip_venv(src_root) -> bool:
+    return (
+        os.environ.get("PWNDBG_VENV_PATH") == "PWNDBG_PLEASE_SKIP_VENV"
+        or (src_root / ".skip-venv").exists()
+    )
+
+
+class Test:
+    def __init__(self, debugger, _):
+        pass
+
+    def __call__(self, debugger, command, exe_context, result):
+        print(f"{debugger}, {command}, {exe_context}, {result}")
+
+
+def main(debugger: lldb.SBDebugger) -> None:
+    profiler = cProfile.Profile()
+
+    start_time = None
+    if os.environ.get("PWNDBG_PROFILE") == "1":
+        start_time = time.time()
+        profiler.enable()
+
+    src_root = Path(__file__).parent.resolve()
+    if not skip_venv(src_root):
+        venv_path = get_venv_path(src_root)
+        if not venv_path.exists():
+            print(f"Cannot find Pwndbg virtualenv directory: {venv_path}. Please re-run setup.sh")
+            sys.exit(1)
+
+        update_deps(src_root, venv_path)
+        fixup_paths(src_root, venv_path)
+
+    os.environ["PWNLIB_NOTERM"] = "1"
+
+    import pwndbg  # noqa: F811
+    import pwndbg.dbg.lldb
+
+    pwndbg.dbg = pwndbg.dbg_mod.lldb.LLDB()
+    pwndbg.dbg.setup(debugger)
+
+    import pwndbg.lldblib
+
+    pwndbg.lldblib.register_class_as_cmd(debugger, "test", Test)
+
+    import pwndbg.profiling
+
+    pwndbg.profiling.init(profiler, start_time)
+    if os.environ.get("PWNDBG_PROFILE") == "1":
+        pwndbg.profiling.profiler.stop("pwndbg-load.pstats")
+        pwndbg.profiling.profiler.start()
+
+
+def __lldb_init_module(debugger, _):
+    """
+    Actually handles the setup bits for LLDB.
+
+    LLDB, unlike GDB, exposes the bits we're interested in through object
+    instances, and we are initially only passed the instance for the interactive
+    debugger through this function.
+    """
+
+    main(debugger)

pwndbg/dbg/lldb.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import Any
+from typing import Tuple
+
+import lldb
+from typing_extensions import override
+
+import pwndbg
+
+
+class LLDB(pwndbg.dbg_mod.Debugger):
+    @override
+    def setup(self, *args):
+        debugger = args[0]
+        assert (
+            debugger.__class__ is lldb.SBDebugger
+        ), "lldbinit.py should call setup() with an lldb.SBDebugger object"
+
+        self.debugger = debugger
+
+    @override
+    def get_cmd_window_size(self) -> Tuple[int, int]:
+        import pwndbg.ui
+
+        return pwndbg.ui.get_window_size()
+
+    @override
+    def addrsz(self, address: Any) -> str:
+        return "%#16x" % address
+
+    @override
+    def set_python_diagnostics(self, enabled: bool) -> None:
+        pass

pwndbg/lldblib/__init__.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+
+def register_class_as_cmd(debugger, cmd, c):
+    mod = c.__module__
+    name = c.__qualname__
+    name = f"{mod if mod else ''}.{name}"
+
+    print(debugger.HandleCommand(f"command script add -c {name} -s synchronous {cmd}"))