#!/usr/bin/env python3
# wall2_resolver4.py - HOST-SIDE wall#2 resolver v4 (final recipe).
#
# Association proven on r675: the RM keeps, in ONE page, the client's channel
# window records and their memdescs:
#   push window record {VA=0x4480000, SIZE=0x20000}
#   ring window record {VA=0x44a0000, SIZE=0x40000}
#   push memdesc {0x20000, 0x03008000, FIRST_PA, ptr}
#   ring memdesc {0x40000, 0x03008000, FIRST_PA, ptr}
# where FIRST_PA is the PA of the window's FIRST VA page. Verified: the ring
# memdesc's PA == the service's validated ring pa. So the push memdesc's PA ==
# the PA of VA 0x4480000 (the push page the service cannot translate).
#
# Recipe: wait for a NEW ring line -> find the page containing the ring window
# record {0x44a0000, 0x40000} -> require the sibling push window record
# {0x4480000, 0x20000} and the memdescs in the same page -> cross-check the
# ring memdesc PA == ring_pa -> inject 0x4480000 -> push memdesc PA (sudo).
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_resolver4.log"
PUSH_VA = 0x4480000
RING_VA = 0x44a0000
WAIT_RING_S = int(os.environ.get("W2_WAIT_RING", "420"))
WATCH_S = int(os.environ.get("W2_WATCH", "180"))

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):
    cmd = "echo 0x%x,0x%x,0x%x,%d > %s" % (va, pa, ln, ap, PARAM)
    r = subprocess.run(["sudo", "sh", "-c", cmd], capture_output=True, text=True)
    if r.returncode == 0:
        log("INJECT va=0x%x -> pa=0x%x len=0x%x OK" % (va, pa, ln))
        return True
    log("inject FAIL rc=%d %s" % (r.returncode, r.stderr.strip()))
    return False

def find_qpair(a, b):
    """Find qword pairs {a,b}; returns [(page, off)]."""
    img.seek(0, 2); fsz = img.tell()
    pat = struct.pack("<QQ", a, b)
    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
            gpa = off2gpa(pos + i)
            out.append((gpa & ~0xfff, gpa - (gpa & ~0xfff)))
        pos += CH
    return out

def scan_page_descs(page):
    d = rd(page, 0x1000)
    if not d: return []
    out = []
    for o in range(0, 0x1000 - 24, 8):
        sz, fl = struct.unpack_from("<QQ", d, o)
        if fl != 0x03008000: continue
        if sz not in (0x20000, 0x40000): continue
        pa = struct.unpack_from("<Q", d, o + 16)[0]
        if 0x100000000 <= pa < 0x180000000:
            out.append((sz, o, pa))
    return out

def main():
    log("=== resolver4 start ===")
    def ring_lines():
        return re.findall(r"ring rc=0 pa=0x([0-9a-f]+) ap=", dmesg())
    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 = int(cur[-1], 16)
            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()
    LEARNED = [0x101e39000, 0x102a62000]  # RM window-record pages (stable across runs)
    def try_page(page):
        descs = scan_page_descs(page)
        rdesc = [x for x in descs if x[0] == 0x40000]
        pdesc = [x for x in descs if x[0] == 0x20000]
        if not rdesc or not pdesc:
            return False
        if rdesc[0][2] != ring_pa:
            return False
        ok = False
        for (_sz, o, pa) in pdesc:
            key = (page, o)
            if key in injected: continue
            pd = rd(pa, 0x40)
            log("page=0x%x ring_desc@0x%x(pa=0x%x OK) push_desc@0x%x pa=0x%x head=%s"
                % (page, rdesc[0][1], rdesc[0][2], o, pa,
                   pd.hex() if pd else "?"))
            if inject(PUSH_VA, pa, 0x1000):
                injected.add(key); ok = True
        return ok
    log("learned pages: %s" % [hex(p) for p in LEARNED])
    watch = time.time() + WATCH_S
    last_full = 0
    while time.time() < watch:
        hit = False
        for page in LEARNED:
            if try_page(page):
                hit = True
        if not hit and time.time() - last_full > 5:
            last_full = time.time()
            for (page, _off) in find_qpair(RING_VA, 0x40000):
                if try_page(page):
                    if page not in LEARNED:
                        LEARNED.append(page)
                    hit = True
        time.sleep(0.25)
    log("=== resolver4 done: %d injection(s) ===" % len(injected))
    return 0

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