path: root/debug/risc-v-gsp-rm-vm/greeneye/session.py

"""green-eye instrumentation session: hooks, breakpoints, watchpoints, events.

This is the Frida-like core.  It is deliberately transport-agnostic; the RSP
client is the default backend, and a Session can also be constructed around a
plugin-supplied backend later.
"""
from __future__ import annotations

import re
import time
from dataclasses import dataclass, field as _field
from typing import Callable

from .errors import GreenEyeError, RspTimeout, SymbolNotFound
from .memory import Memory
from .record import NullRecorder, Recorder
from .rsp import RspClient, StopReply
from .symbols import Symbol, SymbolDB

#: Registers saved and restored around ``Session.call`` (RISC-V ABI + pc).
RESTORE_REGS = (
    ["ra", "sp", "gp", "tp", "t0", "t1", "t2", "s0", "s1"]
    + [f"a{i}" for i in range(8)]
    + [f"s{i}" for i in range(2, 12)]
    + [f"t{i}" for i in range(3, 7)]
    + ["pc"]
)


@dataclass
class Hit:
    """Everything known about one stop event."""

    kind: str  # breakpoint | leave | watchpoint | step | fault | exit | stop
    addr: int
    regs: dict = _field(default_factory=dict)
    symbol: Symbol | None = None
    retval: int | None = None
    watch: "Watchpoint | None" = None
    stop: StopReply | None = None
    session: "Session | None" = None
    data: dict = _field(default_factory=dict)

    @property
    def pc(self) -> int:
        return self.regs.get("pc", self.addr)

    def arg(self, index: int) -> int:
        """RISC-V ABI argument a0..a7 -> x10..x17."""
        return self.regs.get(f"a{index}", 0)

    def label(self) -> str:
        if self.symbol is not None:
            return self.symbol.label(self.addr)
        return f"{self.addr:#x}"

    def __str__(self) -> str:  # pragma: no cover - cosmetic
        base = f"[{self.kind}] {self.label()}"
        if self.kind in ("breakpoint", "leave"):
            base += f" a0={self.arg(0):#x} a1={self.arg(1):#x}"
        if self.kind == "leave" and self.retval is not None:
            base += f" -> {self.retval:#x}"
        if self.kind == "watchpoint" and self.watch is not None:
            base += f" {self.watch.kind}@{self.watch.addr:#x}"
        return base


@dataclass
class Breakpoint:
    addr: int
    kind: str = "sw"  # sw | hw
    temporary: bool = False
    purpose: str = "user"  # user | leave
    on_enter: Callable[[Hit], None] | None = None
    on_leave: Callable[[Hit], None] | None = None
    symbol: Symbol | None = None
    attached: bool = True
    hits: int = 0
    listener: str | None = None

    @property
    def label(self) -> str:
        return self.symbol.label(self.addr) if self.symbol else f"{self.addr:#x}"


@dataclass
class Watchpoint:
    addr: int
    size: int = 8
    kind: str = "w"  # w | r | a
    callback: Callable[[Hit], None] | None = None
    symbol: Symbol | None = None
    hits: int = 0
    listener: str | None = None

    @property
    def label(self) -> str:
        return self.symbol.label(self.addr) if self.symbol else f"{self.addr:#x}"


class HookHandle:
    def __init__(self, session: "Session", breakpoint: Breakpoint):
        self._session = session
        self._bp = breakpoint

    @property
    def addr(self) -> int:
        return self._bp.addr

    @property
    def hits(self) -> int:
        return self._bp.hits

    @property
    def label(self) -> str:
        return self._bp.label

    def detach(self) -> None:
        self._session.remove_breakpoint(self._bp)


class Session:
    def __init__(
        self,
        *,
        client: RspClient,
        target=None,
        symbols: SymbolDB | None = None,
        record: str | bool | None = None,
        name: str = "session",
        meta: dict | None = None,
    ):
        self.client = client
        self.target = target
        self.symbols = symbols or SymbolDB()
        self.name = name
        self.image = None
        self._owns_target = target is not None
        if isinstance(record, str):
            self.recorder = Recorder(record, name=name, meta=meta or {})
        elif record:
            self.recorder = Recorder.new(name=name, meta=meta or {})
        else:
            self.recorder = NullRecorder()
        self.memory = Memory(client, symbols=self.symbols, recorder=self.recorder)
        self._bps: dict[int, Breakpoint] = {}
        self._bp_seq = 0
        self._bps_by_addr: dict[int, list[Breakpoint]] = {}
        self._watches: dict[tuple[int, int, str], Watchpoint] = {}
        self._watches_by_addr: dict[int, Watchpoint] = {}
        self._leave_pending: dict[int, list[Breakpoint]] = {}
        self._stopped = True
        self._event_hooks: list[tuple[str, Callable[[Hit], None]]] = []
        self.recorder.event("session_start", name=name)

    # -- convenience proxies ----------------------------------------------

    def load_image(
        self,
        image,
        *,
        write_segments: bool = False,
        set_pc: bool = False,
        bias: int = 0,
        load_symbols: bool = True,
        source: str = "elf",
    ):
        """Register a firmware image with the session.

        ``write_segments`` copies every PT_LOAD through the RSP client (for
        targets started without their loader devices); ``set_pc`` points the
        CPU at the (relocated) entry.  With QEMU's ``-kernel`` fast path both
        are already handled by QEMU, so the default only loads symbols.
        """
        self.image = image
        if load_symbols:
            self.symbols.load_elf(image, bias=bias, source=source)
        if write_segments:
            for addr, data in image.flat_load_plan(bias=bias):
                self.client.write_mem(addr, data)
        if set_pc:
            self.client.write_reg("pc", image.relocated_entry(bias=bias))
        self.recorder.event(
            "image_loaded",
            path=getattr(image, "path", None),
            entry=image.entry,
            bias=bias,
            write_segments=write_segments,
        )
        return image

    @property
    def pc(self) -> int:
        return self.client.pc

    def read_regs(self) -> dict[str, int]:
        return self.client.read_registers()

    def snapshot(self, addr: int, size: int, name: str = "region") -> bytes:
        return self.memory.snapshot(addr, size, name=name)

    def note(self, text: str, **fields) -> None:
        self.recorder.event("note", text=text, **fields)

    # -- symbol helpers ----------------------------------------------------

    def resolve(self, target: int | str) -> int:
        """Resolve an address, a symbol name, or a ``symbol+0xNN`` expression."""
        if isinstance(target, int):
            return target
        if isinstance(target, str):
            if target.lower().startswith("0x"):
                return int(target, 16)
            expression = re.fullmatch(
                r"([A-Za-z_][A-Za-z0-9_.]*)\s*([+-])\s*(0x[0-9a-fA-F]+|\d+)", target.strip()
            )
            if expression:
                base = self.symbols.get(expression.group(1))
                if base is None:
                    raise SymbolNotFound(f"symbol {expression.group(1)!r} not found in symbol DB")
                delta = int(expression.group(3), 0)
                return base.addr + delta if expression.group(2) == "+" else base.addr - delta
            sym = self.symbols.get(target)
            if sym is None:
                raise SymbolNotFound(f"symbol {target!r} not found in symbol DB")
            return sym.addr
        raise GreenEyeError(f"cannot resolve target {target!r}")

    def symbol_at(self, addr: int) -> Symbol | None:
        """Return the symbol containing ``addr``, or the nearest one (alias of
        ``symbols.resolve``); ``None`` when the database is empty."""
        return self.symbols.resolve(addr)

    # -- breakpoints / hooks ----------------------------------------------

    def hook(
        self,
        target: int | str,
        on_enter: Callable[[Hit], None] | None = None,
        on_leave: Callable[[Hit], None] | None = None,
        *,
        kind: str = "sw",
        temporary: bool = False,
        symbol: Symbol | None = None,
        listener: str | None = None,
    ) -> HookHandle:
        addr = self.resolve(target)
        sym = symbol or self.symbols.resolve(addr)
        bp = Breakpoint(
            addr=addr,
            kind=kind,
            temporary=temporary,
            on_enter=on_enter,
            on_leave=on_leave,
            symbol=sym,
            listener=listener,
        )
        self.add_breakpoint(bp)
        self.recorder.event(
            "hook",
            addr=addr,
            symbol=sym.name if sym else None,
            kind=kind,
            on_leave=bool(on_leave),
            listener=listener,
        )
        return HookHandle(self, bp)

    def breakpoint(self, target: int | str, callback: Callable[[Hit], None] | None = None, **kw) -> Breakpoint:
        addr = self.resolve(target)
        bp = Breakpoint(addr=addr, on_enter=callback, symbol=self.symbols.resolve(addr), **kw)
        self.add_breakpoint(bp)
        return bp

    def add_breakpoint(self, bp: Breakpoint) -> None:
        self._bp_seq += 1
        self._bps[self._bp_seq] = bp
        self._bps_by_addr.setdefault(bp.addr, []).append(bp)
        if bp.kind == "hw":
            self.client.set_hw_break(bp.addr)
        else:
            self.client.set_sw_break(bp.addr)

    def remove_breakpoint(self, bp: Breakpoint) -> None:
        for key, candidate in list(self._bps.items()):
            if candidate is bp:
                del self._bps[key]
        lst = self._bps_by_addr.get(bp.addr, [])
        if bp in lst:
            lst.remove(bp)
        if not lst:
            self._bps_by_addr.pop(bp.addr, None)
            self._leave_pending.pop(bp.addr, None)
            try:
                if bp.kind == "hw":
                    self.client.clear_hw_break(bp.addr)
                else:
                    self.client.clear_sw_break(bp.addr)
            except Exception:
                pass
        bp.attached = False

    def on_event(self, kind: str, callback: Callable[[Hit], None]) -> None:
        self._event_hooks.append((kind, callback))

    # -- watchpoints -------------------------------------------------------

    def watch(
        self,
        target: int | str,
        size: int = 8,
        kind: str = "w",
        callback: Callable[[Hit], None] | None = None,
        listener: str | None = None,
    ) -> Watchpoint:
        addr = self.resolve(target)
        wp = Watchpoint(
            addr=addr,
            size=size,
            kind=kind,
            callback=callback,
            symbol=self.symbols.resolve(addr),
            listener=listener,
        )
        key = (addr, size, kind)
        if key in self._watches:
            raise GreenEyeError(f"watchpoint {kind}@{addr:#x}/{size} already exists")
        self._watches[key] = wp
        self._watches_by_addr[addr] = wp
        self.client.set_watch(addr, size=size, kind=kind)
        self.recorder.event(
            "watch",
            addr=addr,
            size=size,
            kind=kind,
            symbol=wp.symbol.name if wp.symbol else None,
            listener=listener,
        )
        return wp

    def unwatch(self, wp: Watchpoint) -> None:
        key = (wp.addr, wp.size, wp.kind)
        self._watches.pop(key, None)
        self._watches_by_addr.pop(wp.addr, None)
        try:
            self.client.clear_watch(wp.addr, size=wp.size, kind=wp.kind)
        except Exception:
            pass

    # -- execution ---------------------------------------------------------

    def step(self, count: int = 1) -> list[Hit]:
        hits: list[Hit] = []
        for _ in range(count):
            stop = self.client.step()
            hit = self._dispatch_stop(stop, from_step=True)
            if hit is not None:
                hits.append(hit)
            if stop.is_exit:
                break
        return hits

    # -- Frida-like remote call / run-until -------------------------------

    def _scratch_addr(self, override: int | None = None) -> int:
        if override:
            return override
        if self.target is not None and getattr(self.target, "memory_mb", 0):
            base = getattr(self.target, "ram_base", 0x80000000)
            return base + self.target.memory_mb * 1024 * 1024 - 0x1000
        if self.image is not None and getattr(self.image, "max_vaddr", 0):
            return self.image.max_vaddr + 0x10000
        return 0x801FF000

    def run_until(self, target: int | str, timeout: float = 5.0) -> Hit:
        """Run until execution reaches ``target`` (Frida's run-to-cursor).

        Sets a temporary breakpoint, services intervening stops, and returns
        the Hit for the target.  Raises ``RspTimeout`` if not reached.
        """
        addr = self.resolve(target)
        if self.pc == addr:
            return Hit(kind="stop", addr=addr, regs=self.read_regs(), symbol=self.symbols.resolve(addr), session=self)
        temp = Breakpoint(addr=addr, kind="sw", temporary=True, on_enter=None)
        self.add_breakpoint(temp)
        deadline = time.monotonic() + timeout
        try:
            while True:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise RspTimeout(f"run_until {addr:#x}: not reached in {timeout}s")
                try:
                    stop = self.client.continue_(timeout=remaining)
                except RspTimeout:
                    self.client.interrupt()
                    raise RspTimeout(f"run_until {addr:#x}: not reached in {timeout}s")
                hit = self._dispatch_stop(stop)
                if hit is not None and hit.addr == addr and hit.kind == "breakpoint":
                    return hit
                if stop.is_exit:
                    raise GreenEyeError(f"run_until {addr:#x}: target exited")
        finally:
            self.remove_breakpoint(temp)

    def call(self, target: int | str, args: tuple | list = (), timeout: float = 5.0, scratch: int | None = None) -> int:
        """Call a function in the target from the debugger (Frida NativeFunction).

        Writes an ``ebreak`` trampoline into a scratch page, sets ``ra`` to it,
        loads ``a0..a7`` and a private stack, jumps to ``target``, services
        intervening stops (hooks/watchpoints included), and returns ``a0``.
        All caller-saved registers and ``pc`` are restored afterwards.
        """
        addr = self.resolve(target)
        scratch = self._scratch_addr(scratch)
        stack = scratch + 0x800
        saved = self.read_regs()
        temp = Breakpoint(addr=scratch, kind="sw", temporary=True, purpose="user")
        self.add_breakpoint(temp)
        try:
            for name in ("a0", "a1", "a2", "a3", "a4", "a5", "a6", "a7"):
                self.client.write_reg(name, 0)
            for index, value in enumerate(args[:8]):
                self.client.write_reg(f"a{index}", int(value) & 0xFFFF_FFFF_FFFF_FFFF)
            self.client.write_reg("ra", scratch)
            self.client.write_reg("sp", stack)
            self.client.write_reg("pc", addr)
            self.recorder.event("call_enter", addr=addr, args=list(args), scratch=scratch)
            deadline = time.monotonic() + timeout
            while True:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise RspTimeout(f"call {addr:#x}: no return in {timeout}s")
                try:
                    stop = self.client.continue_(timeout=remaining)
                except RspTimeout:
                    self.client.interrupt()
                    raise RspTimeout(f"call {addr:#x}: no return in {timeout}s")
                hit = self._dispatch_stop(stop)
                pc = self.client.pc
                if pc == scratch or (hit is not None and hit.addr == scratch):
                    retval = self.client.read_reg("a0")
                    self.recorder.event("call_return", addr=addr, retval=retval)
                    return retval
                if stop.is_exit:
                    raise GreenEyeError(f"call {addr:#x}: target exited")
        finally:
            if temp.attached:
                self.remove_breakpoint(temp)
            for name in RESTORE_REGS:
                if name in saved:
                    try:
                        self.client.write_reg(name, saved[name])
                    except Exception:
                        pass

    def run(
        self,
        duration: float | None = None,
        max_stops: int | None = None,
        timeout: float | None = None,
    ) -> int:
        """Resume and service events until the deadline/budget/exit.

        ``duration`` (or ``timeout`` when duration is None) is a wall-clock
        budget: when it expires the target is interrupted once, the resulting
        stop is dispatched, and the call returns.  Returns the number of stops
        serviced.
        """
        limit = duration if duration is not None else timeout
        deadline = time.monotonic() + limit if limit is not None else None
        stops = 0
        while True:
            if deadline is not None:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    break  # budget spent while stopped: leave the target stopped
                remaining = max(0.001, remaining)
            else:
                remaining = None
            interrupted = False
            try:
                stop = self.client.continue_(timeout=remaining)
            except RspTimeout:
                if deadline is None:
                    raise
                stop = self.client.interrupt()
                interrupted = True
            stops += 1
            self._dispatch_stop(stop)
            if stop.is_exit:
                break
            if max_stops is not None and stops >= max_stops:
                break
            if interrupted or (deadline is not None and time.monotonic() >= deadline):
                break
        return stops

    def stop(self) -> Hit | None:
        stop = self.client.interrupt()
        return self._dispatch_stop(stop)

    # -- stop dispatch -----------------------------------------------------

    def _dispatch_stop(self, stop: StopReply, from_step: bool = False) -> Hit | None:
        if stop.is_exit:
            hit = Hit(kind="exit", addr=0, stop=stop, session=self)
            self.recorder.event("exit", status=stop.exit_status)
            self._fire_event_hooks(hit)
            return hit

        regs = self.client.read_registers()
        pc = regs.get("pc", 0)
        sym = self.symbols.resolve(pc)

        if stop.watch_addr is not None:
            wp = self._watches_by_addr.get(stop.watch_addr)
            hit = Hit(
                kind="watchpoint",
                addr=stop.watch_addr,
                regs=regs,
                symbol=self.symbols.resolve(stop.watch_addr),
                watch=wp,
                stop=stop,
                session=self,
            )
            if wp is not None:
                wp.hits += 1
                if wp.callback is not None:
                    wp.callback(hit)
            self.recorder.event(
                "watch_hit",
                addr=stop.watch_addr,
                kind=wp.kind if wp else stop.watch_kind,
                size=wp.size if wp else None,
                symbol=hit.symbol.name if hit.symbol else None,
                pc=pc,
                regs=_json_regs(regs),
            )
            self._fire_event_hooks(hit)
            # QEMU stops *at* the instruction that performed the watched
            # access.  Step it once with the watchpoint temporarily cleared so
            # resuming does not re-trigger on the same store forever.
            try:
                if wp is not None:
                    self.client.clear_watch(wp.addr, size=wp.size, kind=wp.kind)
                self.client.step()
            except RspTimeout:
                pass
            finally:
                if wp is not None:
                    try:
                        self.client.set_watch(wp.addr, size=wp.size, kind=wp.kind)
                    except Exception:
                        pass
            return hit

        hit = Hit(
            kind="step" if from_step else "stop",
            addr=pc,
            regs=regs,
            symbol=sym,
            stop=stop,
            session=self,
        )

        bps = list(self._bps_by_addr.get(pc, []))
        leave_queue = self._leave_pending.get(pc, [])
        if not bps and not leave_queue:
            self.recorder.event(
                "stop",
                pc=pc,
                symbol=sym.name if sym else None,
                signal=stop.signal,
                reason="step" if from_step else "halt",
            )
            self._fire_event_hooks(hit)
            return hit

        hit.kind = "breakpoint"
        self.recorder.event(
            "bp_hit",
            pc=pc,
            addr=pc,
            symbol=sym.name if sym else None,
            signal=stop.signal,
            regs=_json_regs(regs),
        )
        for bp in bps:
            bp.hits += 1
            if bp.on_enter is not None:
                bp.on_enter(hit)
            if bp.on_leave is not None and not bp.temporary:
                self._arm_leave(regs.get("ra", 0), bp)
            if bp.temporary:
                self.remove_breakpoint(bp)

        # Return hooks: pop one pending hook for this return address.
        if leave_queue:
            bp = leave_queue.pop()
            if not leave_queue:
                self._leave_pending.pop(pc, None)
                for candidate in list(self._bps_by_addr.get(pc, [])):
                    if candidate.purpose == "leave":
                        self.remove_breakpoint(candidate)
            hit.kind = "leave"
            hit.retval = regs.get("a0", 0)
            if bp.symbol is not None:
                hit.symbol = bp.symbol
                hit.addr = bp.addr
            if bp.on_leave is not None:
                bp.on_leave(hit)

        # Software breakpoints need a single-step over them before resuming.
        self._step_over(pc)
        self._fire_event_hooks(hit)
        return hit

    def _arm_leave(self, retaddr: int, bp: Breakpoint) -> None:
        if not retaddr:
            return
        queue = self._leave_pending.get(retaddr)
        if queue is None:
            queue = []
            self._leave_pending[retaddr] = queue
            temp = Breakpoint(
                addr=retaddr,
                kind="sw",
                temporary=True,
                purpose="leave",
                symbol=self.symbols.resolve(retaddr),
            )
            self.add_breakpoint(temp)
        queue.append(bp)

    def _step_over(self, pc: int) -> None:
        sw = [bp for bp in self._bps_by_addr.get(pc, []) if bp.kind == "sw"]
        if not sw:
            return
        for bp in sw:
            self.client.clear_sw_break(pc)
        self.client.step()
        for bp in sw:
            if bp.attached:
                self.client.set_sw_break(pc)

    def _fire_event_hooks(self, hit: Hit) -> None:
        for kind, callback in self._event_hooks:
            if kind in (hit.kind, "*"):
                callback(hit)

    # -- lifecycle ---------------------------------------------------------

    def close(self, kill_target: bool | None = None) -> None:
        try:
            for bp in list(self._bps.values()):
                self.remove_breakpoint(bp)
            for wp in list(self._watches.values()):
                self.unwatch(wp)
        except Exception:
            pass
        self.client.close()
        if kill_target is None:
            kill_target = self._owns_target
        if kill_target and self.target is not None and not getattr(self.target, "keep_alive", False):
            self.target.kill()
        self.recorder.close()
        self._owns_target = False

    def __enter__(self) -> "Session":
        return self

    def __exit__(self, *exc) -> None:
        self.close()


def _json_regs(regs: dict[str, int]) -> dict[str, int]:
    return {name: value for name, value in regs.items() if isinstance(value, int)}