#!/usr/bin/env python3
# live_walk.py - live PT-walk + inject loop for the running client.
# Polls dmesg for the published page directory (pdb) and for the VAs the
# service is failing to read (sema-entry: unread / entry[] lines / registered
# gpfifo), walks each VA through the pdb with the service's own GMMU math,
# and injects resolved pairs via sema_map_add (ap=1 sysmem).
import mmap, re, subprocess, sys, time

F = "/dev/shm/vgpu-guest-ram"
PARAM = "/sys/module/vgsp/parameters/sema_map_add"
SEMA_PA_MAX = 0x180000000

WINDOW_VAS = [0x44a0000, 0x4500000, 0x4560000,      # rings
              0x4480000, 0x4485000, 0x4ca0000, 0x4ca5000,
              0x4d00000, 0x4b17000, 0x4020000, 0x4100000, 0x4160000,
              0x4213000, 0x45a0bb0,                  # ctx-init releases (deterministic)
              0x124935060, 0x102e34060, 0x108d4f060, 0x12a3a4060, 0x116c0b060,  # notif families
              0x20400f000, 0x20400ffc0, 0x20480fef0, 0x204827f7c]  # counter pages / tracker VAs

mm = None
def off(pa): return pa if pa < 0x80000000 else pa - 0x80000000
def rd64(pa):
    if pa + 8 > SEMA_PA_MAX: return None
    d = mm[off(pa):off(pa)+8]
    return int.from_bytes(d, "little") if len(d) == 8 else None
def pde_ap(ap): return "FB" if ap == 1 else ("SYS" if ap in (2, 3) else None)
def pte_ap(ap): return "FB" if ap == 0 else (None if ap == 1 else "SYS")
def pte_addr(e, fb): return ((e >> 8) & 0x1FFFFFF) << 12 if fb else ((e >> 8) & 0x3FFFFFFFFFFF) << 12
def pte_leaf(table, idx):
    e = rd64(table + 8 * idx)
    if e is None or not (e & 1): return None
    ap = pte_ap((e >> 1) & 3)
    if ap is None: return None
    a = pte_addr(e, ap == "FB")
    return (a, ap) if a + 8 <= SEMA_PA_MAX else None

def walk(pdb, va, entries=4):
    import math
    root_bits = max(1, min(12, int(math.log2(entries or 1))))
    table = pdb
    for level in range(3):
        if level == 0:
            bits, shift = root_bits, 49 - root_bits
        else:
            bits, shift = 9, 38 if level == 1 else 29
        idx = (va >> shift) & ((1 << bits) - 1)
        e = rd64(table + 8 * idx)
        if e is None or (e & 1): return None
        ap = pde_ap((e >> 1) & 3)
        if ap is None: return None
        table = pte_addr(e, ap == "FB")
    idx = (va >> 21) & 0xff
    big = rd64(table + 16 * idx); small = rd64(table + 16 * idx + 8)
    if big is None or small is None: return None
    if not (small & 1):
        sap = pde_ap((small >> 1) & 3)
        if sap is not None:
            r = pte_leaf(pte_addr(small, sap == "FB"), (va >> 12) & 0x1ff)
            if r: return r[0] + (va & 0xfff), r[1]
    if big & 1:
        bap = pte_ap((big >> 1) & 3)
        if bap is not None:
            a = pte_addr(big, bap == "FB") + (va & ((1 << 21) - 1))
            if a + 8 <= SEMA_PA_MAX: return a, bap
    else:
        bap = pde_ap((big >> 1) & 3)
        if bap is not None:
            r = pte_leaf(pte_addr(big, bap == "FB"), (va >> 16) & 0x1f)
            if r: return r[0] + (va & 0xffff), r[1]
    return None

def inject(va, pa, ln=0x1000):
    cmd = "echo 0x%x,0x%x,0x%x,1 > %s" % (va & ~0xfff, pa & ~0xfff, ln, PARAM)
    r = subprocess.run(["sudo", "sh", "-c", cmd], capture_output=True)
    print("  INJECT 0x%x -> 0x%x len=0x%x : %s" %
          (va & ~0xfff, pa & ~0xfff, ln, "ok" if r.returncode == 0 else "FAIL"), flush=True)
    return r.returncode == 0

def main():
    global mm
    f = open(F, "rb")
    mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    dur = float(sys.argv[1]) if len(sys.argv) > 1 else 300
    end = time.time() + dur
    pdbs = []          # (pdb_pa, entries)
    done = set()       # (pdb, va_page)
    seen_va = set()
    last_lines = 0
    print("# live_walk start (%.0fs)" % dur, flush=True)
    while time.time() < end:
        out = subprocess.run(["sudo", "dmesg"], capture_output=True).stdout.decode(errors="ignore")
        lines = out.splitlines()[-500:]
        out = "\n".join(lines)
        # pdb offers/registrations
        for m in re.finditer(r"page directory (?:offered as candidate )?pdb=0x([0-9a-f]+) entries=(\d+)", out):
            pdb, ent = int(m.group(1), 16), int(m.group(2))
            if (pdb, ent) not in pdbs:
                pdbs.append((pdb, ent))
                print("# pdb 0x%x entries=%u" % (pdb, ent), flush=True)
        # VAs the service is failing on
        vas = []
        for m in re.finditer(r"sema-entry: unread h=0x[0-9a-f]+ idx=\d+ va=0x([0-9a-f]+)", out):
            vas.append(int(m.group(1), 16))
        for m in re.finditer(r"sema-chan: entry\[\d+\] h=0x[0-9a-f]+ eng=\d+ raw=[0-9a-f]+ va=0x([0-9a-f]+)", out):
            vas.append(int(m.group(1), 16))
        for m in re.finditer(r"registered gpfifo=0x([0-9a-f]+)", out):
            vas.append(int(m.group(1), 16))
        for m in re.finditer(r"GR ring h=0x[0-9a-f]+ rank=\d+ not at fb 0x[0-9a-f]+ yet \(e0=0x[0-9a-f]+ e1=0x[0-9a-f]+ va=0x([0-9a-f]+)", out):
            vas.append(int(m.group(1), 16))
        for m in re.finditer(r"rel=0x([0-9a-f]+) ", out):
            vas.append(int(m.group(1), 16))
        for m in re.finditer(r"notif=0x([0-9a-f]+)", out):
            vas.append(int(m.group(1), 16))
        for m in re.finditer(r"counter page 0x([0-9a-f]+) unreadable", out):
            vas.append(int(m.group(1), 16))
        for m in re.finditer(r"release (?:learned|derived) h=0x[0-9a-f]+ pa=0x[0-9a-f]+ \(from release va=0x([0-9a-f]+)", out):
            vas.append(int(m.group(1), 16))
        vas = [v for v in vas if 0x4000000 <= v < 0x2000000000]
        for va in WINDOW_VAS + vas:
            seen_va.add(va)
        # resolve + inject through every pdb
        for (pdb, ent) in pdbs:
            for va in list(seen_va):
                key = (pdb, va & ~0xfff)
                if key in done:
                    continue
                r = walk(pdb, va, ent)
                if not r:
                    continue
                done.add(key)
                pa, ap = r
                if ap != "SYS":
                    print("# va=0x%x -> pa=0x%x ap=%s (skip non-sys)" % (va, pa, ap), flush=True)
                    continue
                print("# walk ok pdb=0x%x va=0x%x -> pa=0x%x" % (pdb, va, pa), flush=True)
                inject(va, pa)
        time.sleep(0.4)
    mm.close(); f.close()
    print("# live_walk done", flush=True)

if __name__ == "__main__":
    main()