#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
#
# cudatool.py - Open-IOV CUDA userspace driver analysis toolkit.
#
# A dependency-free (stdlib-only) ELF toolkit for the NVIDIA CUDA userspace
# driver, usable on every driver version without a Ghidra project. It gives
# the fast facts that do not require decompilation: exports, function ranges
# derived from the exception-handling tables, the CUDA forward-compatibility
# note, and a dmesg/crash-address decoder. It is also the version-delta tool:
# `delta` compares two driver versions and reports exactly what changed at
# the interface level, which is what keeps the Open-IOV side anti-fragile
# across driver upgrades.
#
# Commands
#   info <lib>                     ELF/version summary
#   exports <lib> [--json]         exported symbols (file vaddrs)
#   functions <lib> [--json]       function ranges from .eh_frame_hdr
#   nearest <lib> <hex_addr>       nearest preceding function/export for a vaddr
#   crash <lib> <map_base> <ip>    dmesg mapping base + ip -> file vaddr + symbol
#   strings <lib> [--grep RE]      defined strings with file offsets
#   delta <libA> <libB>            interface-level diff between two versions
#   wrappers <lib>                 export stub -> generated wrapper target
#   note <lib>                     CUDA forward-compatibility note (JSON)
#
# Address convention: every address printed by this tool is a FILE VIRTUAL
# ADDRESS, i.e. the `libcuda+0x...` numbers used in the project docs. Ghidra
# may load the same file at a non-zero image base (0x100000 for these builds);
# use tools/ghidra_decompile_at.sh which translates file vaddrs automatically.
#
# License note: this tool only *reads* the vendor binary for interoperability
# analysis; it contains no vendor code and is published under BSD-3-Clause.

import argparse
import json
import re
import struct
import sys

ELF_MAGIC = b"\x7fELF"


class ELFError(Exception):
    pass


class ELF:
    """Minimal 64-bit little-endian ELF reader."""

    def __init__(self, path):
        self.path = path
        with open(path, "rb") as f:
            self.data = f.read()
        if self.data[:4] != ELF_MAGIC:
            raise ELFError(f"{path}: not an ELF file")
        if self.data[4] != 2:
            raise ELFError(f"{path}: not ELF64")
        if self.data[5] != 1:
            raise ELFError(f"{path}: not little-endian")
        (self.e_type, self.e_machine, _v, self.e_entry, self.e_phoff, self.e_shoff,
         self.e_flags, self.e_ehsize, self.e_phentsize, self.e_phnum,
         self.e_shentsize, self.e_shnum, self.e_shstrndx) = struct.unpack_from(
            "<HHIQQQIHHHHHH", self.data, 16)

        self.program_headers = []
        for i in range(self.e_phnum):
            off = self.e_phoff + i * self.e_phentsize
            (p_type, p_flags, p_offset, p_vaddr, p_paddr, p_filesz, p_memsz,
             p_align) = struct.unpack_from("<IIQQQQQQ", self.data, off)
            self.program_headers.append(dict(
                type=p_type, flags=p_flags, offset=p_offset, vaddr=p_vaddr,
                paddr=p_paddr, filesz=p_filesz, memsz=p_memsz, align=p_align,
                index=i))

        self.sections = []
        self.section_by_name = {}
        for i in range(self.e_shnum):
            off = self.e_shoff + i * self.e_shentsize
            (sh_name, sh_type, sh_flags, sh_addr, sh_offset, sh_size,
             sh_link, sh_info, sh_addralign, sh_entsize) = struct.unpack_from(
                "<IIQQQQIIQQ", self.data, off)
            self.sections.append(dict(
                index=i, name_off=sh_name, type=sh_type, flags=sh_flags,
                addr=sh_addr, offset=sh_offset, size=sh_size, link=sh_link,
                info=sh_info, addralign=sh_addralign, entsize=sh_entsize))
        if self.e_shstrndx:
            shstr = self.sections[self.e_shstrndx]
            base = shstr["offset"]
            for s in self.sections:
                end = self.data.find(b"\0", base + s["name_off"])
                s["name"] = self.data[base + s["name_off"]:end].decode(
                    "utf-8", "replace")
                self.section_by_name[s["name"]] = s

    # ---- helpers ---------------------------------------------------------

    def section(self, name):
        return self.section_by_name.get(name)

    def exec_segment_base(self):
        """Base address reported by dmesg/libcuda mappings: first executable
        PT_LOAD vaddr aligned as Linux maps it (page-aligned start)."""
        for ph in self.program_headers:
            if ph["type"] == 1 and (ph["flags"] & 1):  # PT_LOAD | PF_X
                return ph["vaddr"] & ~0xFFF
        raise ELFError("no executable PT_LOAD segment")

    def read_string(self, offset):
        end = self.data.find(b"\0", offset)
        return self.data[offset:end].decode("utf-8", "replace")

    def build_id(self):
        for s in self.sections:
            if s["name"] == ".note.gnu.build-id":
                off = s["offset"]
                namesz, descsz, ntype = struct.unpack_from("<III", self.data, off)
                name = self.data[off + 12:off + 12 + namesz].rstrip(b"\0")
                if name == b"GNU" and ntype == 3:
                    desc = self.data[off + 12 + ((namesz + 3) & ~3):
                                      off + 12 + ((namesz + 3) & ~3) + descsz]
                    return desc.hex()
        return None

    # ---- symbols ---------------------------------------------------------

    def dynamic_symbols(self):
        symtab = self.section(".dynsym")
        if not symtab:
            return []
        strtab = self.sections[symtab["link"]]
        out = []
        n = symtab["size"] // 24
        for i in range(n):
            off = symtab["offset"] + i * 24
            st_name, st_info, st_other, st_shndx, st_value, st_size = \
                struct.unpack_from("<IBBHQQ", self.data, off)
            name = self.read_string(strtab["offset"] + st_name)
            out.append(dict(
                name=name,
                value=st_value,
                size=st_size,
                type=st_info & 0xF,
                bind=st_info >> 4,
                shndx=st_shndx,
                defined=st_shndx != 0,
            ))
        return out

    def exports(self):
        return [s for s in self.dynamic_symbols()
                if s["defined"] and s["type"] == 2 and s["value"]]

    def imports(self):
        return [s for s in self.dynamic_symbols()
                if not s["defined"] and s["name"]]

    # ---- exception handling function table -------------------------------

    def function_ranges(self):
        """Return [(start, end)] from .eh_frame_hdr.

        .eh_frame_hdr layout (DWARF .eh_frame_hdr, version 1):
          u8 version; u8 eh_frame_ptr_enc; u8 fde_count_enc; u8 table_enc;
          encoded eh_frame_ptr; encoded fde_count;
          table[count] { encoded pc_begin; encoded fde_addr; }

        GCC emits: eh_frame_ptr_enc = DW_EH_PE_pcrel|sdata4 (0x1b),
        fde_count_enc = udata4 (0x03), table_enc = DW_EH_PE_datarel|sdata4
        (0x3b). The datarel anchor is the start of .eh_frame_hdr. The parser
        below handles the encodings generally (pcrel/datarel + 1/2/4/8-byte
        and ULEB/SLEB), so it also works for other toolchains.
        """
        hdr = self.section(".eh_frame_hdr")
        if not hdr:
            return []
        base = hdr["offset"]
        anchor = hdr["addr"]
        version = self.data[base]
        if version != 1:
            raise ELFError(f"unsupported .eh_frame_hdr version {version}")
        eh_ptr_enc, fde_count_enc, table_enc = self.data[base + 1:base + 4]

        def enc_size(enc):
            return {0x0: 0, 0x1: 1, 0x2: 2, 0x3: 4, 0x4: 8,
                    0x9: 1, 0xA: 2, 0xB: 4, 0xC: 8}.get(enc & 0x0F, 0)

        def decode(enc, ptr, field_addr):
            fmt = enc & 0x0F
            if fmt == 0x0:
                return 0
            if fmt == 0x1:
                val, shift = 0, 0
                while True:
                    b = self.data[ptr]
                    ptr += 1
                    val |= (b & 0x7F) << shift
                    if not (b & 0x80):
                        break
                    shift += 7
                return val
            if fmt == 0x2:
                return struct.unpack_from("<H", self.data, ptr)[0]
            if fmt == 0x3:
                return struct.unpack_from("<I", self.data, ptr)[0]
            if fmt == 0x4:
                return struct.unpack_from("<Q", self.data, ptr)[0]
            if fmt == 0x9:
                val, shift = 0, 0
                while True:
                    b = self.data[ptr]
                    ptr += 1
                    val |= (b & 0x7F) << shift
                    shift += 7
                    if not (b & 0x80):
                        if b & 0x40:
                            val -= 1 << shift
                        break
                return val
            if fmt == 0xA:
                return struct.unpack_from("<h", self.data, ptr)[0]
            if fmt == 0xB:
                return struct.unpack_from("<i", self.data, ptr)[0]
            if fmt == 0xC:
                return struct.unpack_from("<q", self.data, ptr)[0]
            raise ELFError(f"unsupported eh encoding 0x{enc:02x}")

        def absolute(enc, ptr):
            val = decode(enc, ptr, anchor)
            kind = enc & 0x70
            if kind == 0x10:      # DW_EH_PE_pcrel: relative to field address
                val += anchor + (ptr - base)
            elif kind == 0x30:    # DW_EH_PE_datarel
                val += anchor
            return val

        p = base + 4
        p += enc_size(eh_ptr_enc)          # skip eh_frame_ptr
        count = decode(fde_count_enc, p, anchor)
        p += enc_size(fde_count_enc)
        starts = []
        for _ in range(count):
            pc = absolute(table_enc, p)
            p += enc_size(table_enc)
            p += enc_size(table_enc)       # skip fde address
            starts.append(pc)
        starts.sort()
        out = []
        for i, pc in enumerate(starts):
            end = starts[i + 1] if i + 1 < len(starts) else None
            out.append((pc, end))
        return out

    # ---- CUDA forward compatibility note ---------------------------------

    def cuda_note(self):
        s = self.section(".note.cuda.fwd_compatibility")
        if not s:
            return None
        off = s["offset"]
        namesz, descsz, ntype = struct.unpack_from("<III", self.data, off)
        name_off = off + 12
        name = self.data[name_off:name_off + namesz].rstrip(b"\0").decode()
        desc_off = name_off + ((namesz + 3) & ~3)
        desc = self.data[desc_off:desc_off + descsz].rstrip(b"\0").decode()
        try:
            parsed = json.loads(desc)
        except json.JSONDecodeError:
            parsed = None
        return dict(owner=name, raw=desc, parsed=parsed)

    def debuglink(self):
        s = self.section(".gnu_debuglink")
        if not s:
            return None
        name = self.read_string(s["offset"])
        crc = struct.unpack_from("<I", self.data, s["offset"] + (len(name) + 4) // 4 * 4)[0]
        return dict(name=name, crc=crc)

    # ---- export stubs -> generated wrapper ---------------------------------

    def relative_relocs(self):
        """Map address -> R_X86_64_RELATIVE addend (the initial slot value)."""
        out = {}
        s = self.section(".rela.dyn")
        if not s:
            return out
        n = s["size"] // 24
        for i in range(n):
            off = s["offset"] + i * 24
            r_offset, r_info, r_addend = struct.unpack_from("<QQq", self.data, off)
            if (r_info & 0xFFFFFFFF) == 8:  # R_X86_64_RELATIVE
                out[r_offset] = r_addend
        return out

    def wrappers(self):
        """Resolve each exported cu* stub to its initial wrapper target.

        The generated stubs are:
            cmp dword ptr [rip+d], 0x321cba00   (7 bytes)
            je  +0xc                            (2 bytes)
            jmp qword ptr [rip+d]               (6 bytes, disp at stub+14)
        so the interposition slot is stub + 18 + disp32. Its initial value is
        set by a R_X86_64_RELATIVE relocation and points at the generated
        wrapper (tracing/tool dispatch), whose first callee is the hand-written
        implementation.
        """
        relocs = self.relative_relocs()
        out = []
        for s in self.exports():
            if not s["name"].startswith("cu"):
                continue
            a = s["value"]
            # sanity-check the cmp/jmp pattern before trusting the arithmetic
            if self.data[a] != 0x81 or self.data[a + 10] != 0x74 or \
                    self.data[a + 12] != 0xFF or self.data[a + 13] != 0x25:
                continue
            disp = struct.unpack_from("<i", self.data, a + 14)[0]
            slot = a + 18 + disp
            target = relocs.get(slot)
            out.append(dict(name=s["name"], stub=a, slot=slot, wrapper=target))
        return out

    def text_bounds(self):
        t = self.section(".text")
        if not t:
            return None
        return (t["addr"], t["addr"] + t["size"])


# ---- commands ------------------------------------------------------------


def cmd_info(args):
    elf = ELF(args.lib)
    tb = elf.text_bounds()
    note = elf.cuda_note()
    print(f"file            : {args.lib}")
    print(f"build id        : {elf.build_id()}")
    print(f"e_type          : {elf.e_type} (3=DYN shared object)")
    print(f"exec segment base: 0x{elf.exec_segment_base():x}")
    print(f".text           : 0x{tb[0]:x} .. 0x{tb[1]:x} ({tb[1]-tb[0]} bytes)")
    print(f"exports         : {len(elf.exports())}")
    print(f"imports         : {len(elf.imports())}")
    if note and note["parsed"]:
        p = note["parsed"]
        print(f"cuda note       : CUDA {p.get('CUDA version')}, "
              f"driver ABI {p.get('driver')}, device {p.get('device')}")
    dl = elf.debuglink()
    if dl:
        print(f"debuglink       : {dl['name']} (crc 0x{dl['crc']:08x})")
    for s in elf.sections:
        print(f"  section {s['name']:<28} addr=0x{s['addr']:08x} "
              f"off=0x{s['offset']:08x} size={s['size']}")


def cmd_exports(args):
    elf = ELF(args.lib)
    ex = sorted(elf.exports(), key=lambda s: s["value"])
    if args.json:
        print(json.dumps([{k: v for k, v in s.items()} for s in ex], indent=1))
        return
    for s in ex:
        print(f"0x{s['value']:08x} {s['size']:>8} {s['name']}")


def cmd_functions(args):
    elf = ELF(args.lib)
    fr = elf.function_ranges()
    if args.json:
        print(json.dumps([dict(start=hex(a), end=hex(b) if b else None)
                          for a, b in fr], indent=1))
        return
    for a, b in fr:
        print(f"0x{a:08x} 0x{b:08x}" if b else f"0x{a:08x} -")


def cmd_nearest(args):
    elf = ELF(args.lib)
    target = int(args.addr, 16)
    ex = sorted(elf.exports(), key=lambda s: s["value"])
    prev = None
    for s in ex:
        if s["value"] <= target:
            prev = s
        else:
            break
    fr = elf.function_ranges()
    fstart = None
    for a, b in fr:
        if a <= target and (b is None or target < b):
            fstart = a
            break
    print(f"vaddr            : 0x{target:x}")
    print(f"function start   : 0x{fstart:x}" if fstart is not None else
          "function start   : <not covered by .eh_frame_hdr>")
    if prev:
        print(f"nearest export   : {prev['name']} + 0x{target - prev['value']:x} "
              f"(export at 0x{prev['value']:x})")
    else:
        print("nearest export   : <none>")


def cmd_crash(args):
    elf = ELF(args.lib)
    base = int(args.map_base, 16)
    ip = int(args.ip, 16)
    vaddr = ip - base + elf.exec_segment_base()
    print(f"file vaddr       : 0x{vaddr:x}")
    args.addr = hex(vaddr)
    cmd_nearest(args)


def cmd_strings(args):
    elf = ELF(args.lib)
    pat = re.compile(args.grep) if args.grep else None
    rodata = [s for s in elf.sections if s["type"] == 1 and not (s["flags"] & 4)]
    seen = set()
    for s in rodata:
        blob = elf.data[s["offset"]:s["offset"] + s["size"]]
        for m in re.finditer(rb"[\x20-\x7e]{5,}", blob):
            text = m.group().decode()
            if pat and not pat.search(text):
                continue
            key = (s["name"], m.start())
            if key in seen:
                continue
            seen.add(key)
            print(f"0x{s['addr'] + m.start():08x} {text}")


def cmd_delta(args):
    a, b = ELF(args.libA), ELF(args.libB)
    ea = {s["name"]: s for s in a.exports()}
    eb = {s["name"]: s for s in b.exports()}
    added = sorted(set(eb) - set(ea))
    removed = sorted(set(ea) - set(eb))
    moved = []
    resized = []
    for n in sorted(set(ea) & set(eb)):
        if ea[n]["value"] != eb[n]["value"]:
            moved.append((n, ea[n]["value"], eb[n]["value"]))
        if ea[n]["size"] != eb[n]["size"]:
            resized.append((n, ea[n]["size"], eb[n]["size"]))
    print(f"A: {args.libA}")
    print(f"   build-id {a.build_id()} exports {len(ea)}")
    print(f"B: {args.libB}")
    print(f"   build-id {b.build_id()} exports {len(eb)}")
    na, nb = a.cuda_note(), b.cuda_note()
    print(f"note A: {na['raw'] if na else None}")
    print(f"note B: {nb['raw'] if nb else None}")
    tb_a, tb_b = a.text_bounds(), b.text_bounds()
    print(f"text A: 0x{tb_a[0]:x}..0x{tb_a[1]:x} ({tb_a[1]-tb_a[0]})")
    print(f"text B: 0x{tb_b[0]:x}..0x{tb_b[1]:x} ({tb_b[1]-tb_b[0]})")
    print(f"exports added   ({len(added)}): " + ", ".join(added[:40]) +
          (" ..." if len(added) > 40 else ""))
    print(f"exports removed ({len(removed)}): " + ", ".join(removed[:40]) +
          (" ..." if len(removed) > 40 else ""))
    print(f"exports moved   ({len(moved)})" +
          (": " + ", ".join(f"{n} 0x{x:x}->0x{y:x}" for n, x, y in moved[:10])
           if moved else ""))
    print(f"exports resized ({len(resized)})" +
          (": " + ", ".join(f"{n} {x}->{y}" for n, x, y in resized[:10])
           if resized else ""))
    print(f"imports A={len(a.imports())} B={len(b.imports())}")


def cmd_wrappers(args):
    elf = ELF(args.lib)
    ws = elf.wrappers()
    if args.json:
        print(json.dumps(ws, indent=1))
        return
    for w in ws:
        if w["wrapper"] is None:
            continue
        print(f"0x{w['stub']:08x} -> 0x{w['wrapper']:08x} {w['name']}")


def cmd_note(args):
    elf = ELF(args.lib)
    print(json.dumps(elf.cuda_note(), indent=2))


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("info"); p.add_argument("lib"); p.set_defaults(fn=cmd_info)
    p = sub.add_parser("exports"); p.add_argument("lib")
    p.add_argument("--json", action="store_true"); p.set_defaults(fn=cmd_exports)
    p = sub.add_parser("functions"); p.add_argument("lib")
    p.add_argument("--json", action="store_true"); p.set_defaults(fn=cmd_functions)
    p = sub.add_parser("nearest"); p.add_argument("lib"); p.add_argument("addr")
    p.set_defaults(fn=cmd_nearest)
    p = sub.add_parser("crash"); p.add_argument("lib")
    p.add_argument("map_base"); p.add_argument("ip"); p.set_defaults(fn=cmd_crash)
    p = sub.add_parser("strings"); p.add_argument("lib")
    p.add_argument("--grep"); p.set_defaults(fn=cmd_strings)
    p = sub.add_parser("delta"); p.add_argument("libA"); p.add_argument("libB")
    p.set_defaults(fn=cmd_delta)
    p = sub.add_parser("note"); p.add_argument("lib"); p.set_defaults(fn=cmd_note)
    p = sub.add_parser("wrappers"); p.add_argument("lib")
    p.add_argument("--json", action="store_true"); p.set_defaults(fn=cmd_wrappers)

    args = ap.parse_args()
    try:
        args.fn(args)
    except ELFError as e:
        print(f"error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()