#!/usr/bin/env python3
# wall2_resolver3.py - HOST-SIDE wall#2 resolver v3: the RM window-memdesc recipe.
#
# Proven structure (r674/r675): the RM keeps per-window records
#   {..., VA, SIZE, ...}  (head/tail list nodes)
# and per-window memdescs in the same page:
#   {SIZE, 0x03008000, FIRST_PA, ptr}
# where FIRST_PA is the PA of the window's FIRST VA page. Verified: the ring
# window {VA=0x44a0000, SIZE=0x40000} -> memdesc PA == the service's validated
# ring pa (0x1158ce000 r675). Therefore the push window
# {VA=0x4480000, SIZE=0x20000} -> memdesc PA == PA of VA 0x4480000 (the push
# page the sema service cannot translate).
#
# Recipe: wait for the ring line -> find the ring window memdesc
# {0x40000, 0x03008000, RING_PA} -> in the same page find the push window
# memdesc {0x20000, 0x03008000, PUSH_PA} -> inject 0x4480000 -> PUSH_PA.
import subprocess, struct, time, sys, re, os

RAM = "/dev/shm/vgpu-guest-ram"
PARAM = "/sys/module/vgsp/parameters/sema_map_add"
LOG = "/tmp/opencode/wall2_resolver3.log"
PUSH_VA = 0x4480000
WAIT_RING_S = int(os.environ.get("W2_WAIT_RING", "420"))
WATCH_S = int(os.environ.get("W2_WATCH", "150"))

logf = open(LOG, "a", buffering=1)
def log(*a):
    s = " ".join(str(x) for x in a)
    print(s, flush=True)
    logf.write("[%s] %s\n" % (time.strftime("%H:%M:%S"), s))

img = open(RAM, "rb", buffering=0)
def gpa2off(pa):
    if pa < 0x80000000: return pa
    if pa < 0x180000000: return 0x80000000 + (pa - 0x100000000)
    return None
def off2gpa(off):
    return off if off < 0x80000000 else 0x100000000 + (off - 0x80000000)
def rd(pa, n):
    off = gpa2off(pa)
    if off is None: return None
    img.seek(off); return img.read(n)

def dmesg():
    return subprocess.run(["sudo", "dmesg"], capture_output=True, text=True).stdout

def inject(va, pa, ln=0x1000, ap=1):
    try:
        with open(PARAM, "w") as f:
            f.write("0x%x,0x%x,0x%x,%d\n" % (va, pa, ln, ap))
        log("INJECT va=0x%x -> pa=0x%x len=0x%x" % (va, pa, ln))
        return True
    except OSError as e:
        log("inject FAIL", e)
        return False

def find_memdescs(size):
    """Find {size, 0x03008000, PA} patterns; returns [(page, off, pa)]."""
    img.seek(0, 2); fsz = img.tell()
    pat = struct.pack("<QQ", size, 0x03008000)
    out = []
    CH = 64 << 20
    pos = 0
    while pos < fsz:
        img.seek(pos); d = img.read(min(CH, fsz - pos))
        st = 0
        while True:
            i = d.find(pat, st)
            if i < 0: break
            st = i + 1
            if i + 24 > len(d): continue
            pa = struct.unpack_from("<Q", d, i + 16)[0]
            gpa = off2gpa(pos + i)
            if 0x100000000 <= pa < 0x180000000:
                out.append((gpa & ~0xfff, gpa - (gpa & ~0xfff), pa))
        pos += CH
    return out

def main():
    log("=== resolver3 start ===")
    # count pre-existing ring lines; wait for a NEW one (fresh client kick)
    def ring_lines():
        out = []
        for ln in dmesg().splitlines():
            m = re.search(r"ring rc=0 pa=0x([0-9a-f]+) ap=", ln)
            if m:
                out.append(int(m.group(1), 16))
        return out
    base = ring_lines()
    log("pre-existing ring lines: %d" % len(base))
    ring_pa = None
    deadline = time.time() + WAIT_RING_S
    while time.time() < deadline:
        cur = ring_lines()
        if len(cur) > len(base):
            ring_pa = cur[-1]
            break
        time.sleep(1)
    if not ring_pa:
        log("no NEW ring line; exit"); return 1
    log("NEW RING PA = 0x%x" % ring_pa)

    injected = set()
    watch = time.time() + WATCH_S
    while time.time() < watch:
        # ring window memdesc {0x40000, 0x03008000, ring_pa}
        rd_hits = [x for x in find_memdescs(0x40000) if x[2] == ring_pa]
        if not rd_hits:
            time.sleep(0.5); continue
        for (page, off, _pa) in rd_hits:
            d = rd(page, 0x1000)
            if not d: continue
            # find sibling push memdesc in the same page
            for o in range(0, 0x1000 - 24, 8):
                sz, fl = struct.unpack_from("<QQ", d, o)
                if sz == 0x20000 and fl == 0x03008000:
                    pa = struct.unpack_from("<Q", d, o + 16)[0]
                    if not (0x100000000 <= pa < 0x180000000):
                        continue
                    key = (page, pa)
                    if key in injected: continue
                    log("memdesc page=0x%x ring_desc@0x%x push_desc@0x%x push_pa=0x%x"
                        % (page, off, o, pa))
                    # content sanity: read the push page
                    pd = rd(pa, 0x40)
                    log("   push page head: %s" % (pd.hex() if pd else "?"))
                    if inject(PUSH_VA, pa, 0x1000):
                        injected.add(key)
        if injected and ("injected mapping va=0x%x" % PUSH_VA) in dmesg():
            log("service confirmed; keep watching")
        time.sleep(0.5)
    log("=== resolver3 done: %d injection(s) ===" % len(injected))
    return 0

if __name__ == "__main__":
    sys.exit(main())