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

"""Symbol resolution and address translation for risc-v-gsp-rm-debugger.

Inputs supported (all optional, merged into one database):

* ELF ``.symtab`` symbols (demo firmware);
* the project's Ghidra exports:
  - ``names.tsv``   ``<file_vaddr>\t<name>[\t<source>]`` (comments with #),
  - ``decompiled_index.tsv`` / bucket TSVs
    ``<program_addr>\t<file_addr>\t<size>\t<name>\t<bucket>\t<line>\t<status>``,
  - ``exports.tsv``  ``addr\tfile_addr\tsize\tname``,
* risc-v-gsp-rm-debugger native JSON (``greeneye-symbols.json``).

Address conventions (matching the project docs):
  * Ghidra program address = file vaddr + image base (``0x100000`` for libcuda);
  * runtime address        = file vaddr + runtime load bias;
  * for firmware images linked at their load address the bias is 0.

``SymbolDB`` stores runtime addresses by default; each loader takes a ``bias``
so a relocated firmware or a different image base can be normalised.
"""
from __future__ import annotations

import bisect
import json
import os
from dataclasses import dataclass, field as _field

from .errors import GreenEyeError, SymbolNotFound


@dataclass
class Symbol:
    addr: int
    name: str
    size: int = 0
    source: str = ""
    subsystem: str | None = None
    extras: dict = _field(default_factory=dict)

    @property
    def end(self) -> int:
        return self.addr + self.size if self.size else self.addr

    def contains(self, addr: int) -> bool:
        if self.size:
            return self.addr <= addr < self.addr + self.size
        return addr == self.addr

    def label(self, addr: int | None = None) -> str:
        if addr is None or addr == self.addr or not self.size:
            return self.name
        return f"{self.name}+{addr - self.addr:#x}"

    def to_json(self) -> dict:
        out = {"addr": self.addr, "name": self.name}
        if self.size:
            out["size"] = self.size
        if self.source:
            out["source"] = self.source
        if self.subsystem:
            out["subsystem"] = self.subsystem
        if self.extras:
            out.update(self.extras)
        return out


class SymbolDB:
    def __init__(self) -> None:
        self._symbols: list[Symbol] = []
        self._sorted = False
        self._by_name: dict[str, Symbol] = {}
        self._addr_index: list[int] = []

    # -- loading -----------------------------------------------------------

    def add(self, symbol: Symbol) -> None:
        self._symbols.append(symbol)
        self._sorted = False
        existing = self._by_name.get(symbol.name)
        if existing is None or (symbol.source and not existing.source):
            self._by_name[symbol.name] = symbol

    def load_elf(self, image, bias: int = 0, source: str = "elf", only_code: bool = True) -> int:
        """Load ``.symtab`` symbols from a FirmwareImage/ELF path."""
        from .elf import ElfSymbol, load_elf

        if isinstance(image, str):
            image = load_elf(image)
        count = 0
        for sym in image.symbols:  # type: ignore[union-attr]
            if only_code and not (sym.is_function or sym.is_object):
                continue
            self.add(
                Symbol(
                    addr=sym.value + bias,
                    name=sym.name,
                    size=sym.size,
                    source=source,
                    extras={"type": "func" if sym.is_function else "object"},
                )
            )
            count += 1
        return count

    def load_names_tsv(self, path: str, bias: int = 0) -> int:
        """Load the project's Ghidra ``names.tsv`` (file vaddrs)."""
        count = 0
        with open(path, "r", encoding="utf-8", errors="replace") as fh:
            for line in fh:
                line = line.rstrip("\n")
                if not line or line.startswith("#"):
                    continue
                cols = line.split("\t")
                if len(cols) < 2:
                    continue
                try:
                    addr = int(cols[0], 16)
                except ValueError:
                    continue
                self.add(
                    Symbol(addr=addr + bias, name=cols[1], source=cols[2] if len(cols) > 2 else "ghidra-names")
                )
                count += 1
        return count

    def load_index_tsv(self, path: str, bias: int = 0) -> int:
        """Load ``decompiled_index.tsv``/bucket TSVs/exports.tsv heuristically."""
        count = 0
        with open(path, "r", encoding="utf-8", errors="replace") as fh:
            for line in fh:
                line = line.rstrip("\n")
                if not line or line.startswith("#"):
                    continue
                cols = line.split("\t")
                if cols[0] in ("addr", "caller"):
                    continue
                if len(cols) >= 7:
                    program, file_addr, size, name, bucket, lineno, status = cols[:7]
                    addr = _parse_hex(file_addr)
                    extras = {"program_addr": _parse_hex(program), "bucket": bucket, "line": lineno, "status": status}
                elif len(cols) >= 4:
                    program, file_addr, size, name = cols[:4]
                    addr = _parse_hex(file_addr)
                    extras = {"program_addr": _parse_hex(program)}
                elif len(cols) >= 2:
                    addr = _parse_hex(cols[0])
                    name = cols[1]
                    size = "0"
                    extras = {}
                else:
                    continue
                if addr is None or not name:
                    continue
                try:
                    size_val = int(size, 0)
                except ValueError:
                    size_val = 0
                self.add(
                    Symbol(addr=addr + bias, name=name, size=size_val, source="ghidra-index", extras=extras)
                )
                count += 1
        return count

    def load_callgraph(self, path: str) -> int:
        """Attach caller/callee lists from ``callgraph.tsv`` (program addrs)."""
        by_program: dict[int, Symbol] = {}
        for sym in self._symbols:
            prog = sym.extras.get("program_addr")
            if prog is not None:
                by_program[int(prog)] = sym
        if not by_program:
            return 0
        for sym in self._symbols:
            sym.extras.setdefault("callers", [])
            sym.extras.setdefault("callees", [])
        count = 0
        with open(path, "r", encoding="utf-8", errors="replace") as fh:
            for line in fh:
                line = line.rstrip("\n")
                if not line or line.startswith("#") or line.startswith("caller"):
                    continue
                cols = line.split("\t")
                if len(cols) < 2:
                    continue
                caller, callee = _parse_hex(cols[0]), _parse_hex(cols[1])
                if caller is None or callee is None:
                    continue
                if caller in by_program:
                    by_program[caller].extras["callees"].append(callee)
                    count += 1
                if callee in by_program:
                    by_program[callee].extras["callers"].append(caller)
        return count

    def load_json(self, path: str) -> int:
        """Load a risc-v-gsp-rm-debugger JSON symbol DB; raise GreenEyeError on bad input."""
        try:
            with open(path, "r", encoding="utf-8") as fh:
                payload = json.load(fh)
        except (OSError, json.JSONDecodeError) as exc:
            raise GreenEyeError(f"cannot load symbol DB {path!r}: {exc}") from exc
        if not isinstance(payload, (dict, list)):
            raise GreenEyeError(f"symbol DB {path!r} must be a JSON object or list")
        symbols = payload.get("symbols", []) if isinstance(payload, dict) else payload
        count = 0
        for item in symbols:
            try:
                addr = int(item["addr"], 0) if isinstance(item["addr"], str) else int(item["addr"])
            except (KeyError, TypeError, ValueError):
                continue
            extras = {k: v for k, v in item.items() if k not in ("addr", "name", "size", "source", "subsystem")}
            self.add(
                Symbol(
                    addr=addr,
                    name=item.get("name", f"sub_{addr:x}"),
                    size=int(item.get("size", 0)),
                    source=item.get("source", "json"),
                    subsystem=item.get("subsystem"),
                    extras=extras,
                )
            )
            count += 1
        return count

    def save_json(self, path: str, meta: dict | None = None) -> None:
        self._ensure_sorted()
        payload = {
            "format": "greeneye-symbols/1",
            "count": len(self._symbols),
            "meta": meta or {},
            "symbols": [s.to_json() for s in self._symbols],
        }
        with open(path, "w", encoding="utf-8") as fh:
            json.dump(payload, fh, indent=1, sort_keys=True)

    # -- queries -----------------------------------------------------------

    def _ensure_sorted(self) -> None:
        if self._sorted:
            return
        self._symbols.sort(key=lambda s: (s.addr, -s.size))
        self._addr_index = [s.addr for s in self._symbols]
        self._sorted = True

    def __len__(self) -> int:
        return len(self._symbols)

    def __iter__(self):
        return iter(self._symbols)

    def get(self, name: str) -> Symbol | None:
        return self._by_name.get(name)

    def lookup(self, name: str) -> int:
        sym = self._by_name.get(name)
        if sym is None:
            raise SymbolNotFound(f"symbol {name!r} not found")
        return sym.addr

    def resolve(self, addr: int, max_delta: int = 0x10000) -> Symbol | None:
        """Resolve an address to a containing or nearest preceding symbol."""
        self._ensure_sorted()
        if not self._symbols:
            return None
        pos = bisect.bisect_right(self._addr_index, addr) - 1
        if pos < 0:
            return None
        candidate = self._symbols[pos]
        if candidate.contains(addr):
            return candidate
        if candidate.size and addr < candidate.addr + candidate.size:
            return candidate
        # nearest-within-delta fallback for zero-size symbols (Ghidra names)
        if not candidate.size and 0 <= addr - candidate.addr <= max_delta:
            return candidate
        return None

    def annotate(self, addr: int, max_delta: int = 0x10000) -> str:
        sym = self.resolve(addr, max_delta=max_delta)
        if sym is None:
            return f"{addr:#x} <unknown>"
        return f"{addr:#x} {sym.label(addr)}"

    def nearest(self, addr: int, max_delta: int = 0x10000) -> Symbol | None:
        """Alias of :meth:`resolve` kept for call sites that read better as
        "nearest symbol" (same containing/preceding/max-delta rules)."""
        return self.resolve(addr, max_delta=max_delta)

    def names(self) -> list[str]:
        return sorted(self._by_name)

    def by_subsystem(self, subsystem: str) -> list[Symbol]:
        """Return all symbols tagged with ``subsystem`` (loaders set it when known)."""
        return [s for s in self._symbols if s.subsystem == subsystem]


def _parse_hex(value: str) -> int | None:
    value = value.strip()
    if not value:
        return None
    try:
        return int(value, 16)
    except ValueError:
        try:
            return int(value, 0)
        except ValueError:
            return None


def translation_for(
    *,
    file_vaddr: int,
    program_addr: int | None = None,
    runtime_addr: int | None = None,
) -> dict:
    """Explain the three address spaces for a debug session (docs aid)."""
    out = {"file_vaddr": file_vaddr}
    if program_addr is not None:
        out["image_base"] = program_addr - file_vaddr
    if runtime_addr is not None:
        out["load_bias"] = runtime_addr - file_vaddr
    return out