#!/usr/bin/env python3
# wall2_livecap2.py - capture the GR channel-init push page by its release
# group {0x20038090, hi, lo, payload, continuation} at page offset 0x598.
# The push VA comes from the service's own entry[0] decode in dmesg.
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_livecap2.log"
MASK = ~0xfff & ((1 << 64) - 1)
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 prefix_ok(d):
    off = 0
    while off + 4 <= len(d):
        hdr = struct.unpack_from("<I", d, off)[0]
        if hdr: break
        off += 4
    if off + 4 > len(d): return False
    op = (hdr >> 29) & 7
    if op in (4, 7): return True
    if op not in (1, 3, 5): return False
    return ((hdr >> 16) & 0xfff) != 0

def cont_ok(v):
    op = (v >> 29) & 7
    cnt = (v >> 16) & 0xfff
    return (op in (1, 3, 5) and cnt != 0) or op in (4, 7)

def scan():
    img.seek(0, 2); fsz = img.tell()
    pat = b"\x90\x80\x03\x20"
    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
            g = d[i:i+24]
            hdr, hi, lo, pay, w5, v5 = struct.unpack("<6I", g)
            gpa = off2gpa(pos + i)
            page = gpa & MASK
            off = gpa - page
            out.append((page, off, hi, lo, pay, w5, v5))
        pos += CH
    return out

def main():
    log("=== livecap2 start watch=%ds ===" % WATCH_S)
    injected = set()
    deadline = time.time() + WATCH_S
    seen_pages = {}
    while time.time() < deadline:
        hits = scan()
        for (page, off, hi, lo, pay, w5, v5) in hits:
            if not (0 < pay < (1 << 24)):
                continue
            if not cont_ok(w5):
                continue
            d = rd(page, 0x1000)
            if not d: continue
            pok = prefix_ok(d[:0x100])
            key = (page, off)
            if key in seen_pages:
                continue
            seen_pages[key] = 1
            log("hit page=0x%x off=0x%x sem=0x%x pay=%u w5=0x%08x v5=0x%08x prefix_ok=%s"
                % (page, off, (hi << 32) | lo, pay, w5, v5, pok))
            log("    head: %s" % d[:0x30].hex())
            log("    @0x590: %s" % d[0x590:0x5c0].hex())
            if pok and off in (0x598, 0x590, 0x580, 0x588):
                va = 0x4480000
                if va not in injected:
                    try:
                        with open(PARAM, "w") as f:
                            f.write("0x%x,0x%x,0x%x,1\n" % (va, page, 0x1000))
                        log("    INJECT 0x%x -> 0x%x" % (va, page))
                        injected.add(va)
                    except OSError as e:
                        log("    inject FAIL", e)
        time.sleep(0.5)
    log("=== livecap2 done: %d injections ===" % len(injected))
    return 0

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