#!/usr/bin/env python3
# host_resolve.py - host-side push resolver: read the service's sema-pushall
# lines (pa=0x0, non-zero first8), find the content page in
# /dev/shm/vgpu-guest-ram by the push's first8 bytes, and inject the page
# mapping via sema_map_add. Replaces the guest-side kcore walker: the
# service already tells us push VA + content prefix, the host RAM file has
# the bytes.
import mmap, os, re, subprocess, sys, time

RAM = "/dev/shm/vgpu-guest-ram"
PARAM = "/sys/module/vgsp/parameters/sema_map_add"
BASE = 0x80000000  # file offset = pa - BASE

def inject(va_page, pa):
    cmd = "echo 0x%x,0x%x,0x1000,1 > %s" % (va_page, pa, PARAM)
    r = subprocess.run(["sudo", "sh", "-c", cmd], capture_output=True)
    ok = r.returncode == 0
    print("inject 0x%x -> 0x%x : %s" % (va_page, pa, "ok" if ok else r.stderr.decode().strip()))
    return ok

def main():
    seen = set()
    f = open(RAM, "rb")
    mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    injected = 0
    deadline = time.time() + (float(sys.argv[1]) if len(sys.argv) > 1 else 600)
    while time.time() < deadline and injected < 40:
        out = subprocess.run(["sudo", "dmesg"], capture_output=True).stdout.decode(errors="ignore")
        lines = [l for l in out.splitlines() if "sema-pushall" in l]
        for l in lines[-800:]:
            m = re.search(r"va=(0x[0-9a-f]+) pa=0x0 len=\d+ allow=\d trusted=\d first8=([0-9a-f]{8}) ([0-9a-f]{8})", l)
            if not m:
                continue
            va = int(m.group(1), 16)
            w0, w1 = int(m.group(2), 16), int(m.group(3), 16)
            if (w0, w1) == (0, 0) and (w0 == 0 or w1 == 0):
                pass
            if w0 == 0 and w1 == 0:
                continue
            key = (va & ~0xFFF, w0, w1)
            if key in seen:
                continue
            pat = w0.to_bytes(4, "little") + w1.to_bytes(4, "little")
            inpage = va & 0xFFF
            # search with the in-page offset constraint
            start = 0
            hit = None
            while True:
                i = mm.find(pat, start)
                if i < 0:
                    break
                if (i & 0xFFF) == inpage:
                    hit = i
                    break
                start = i + 1
            if hit is None:
                # fallback: any hit within the same page offset modulo
                start = 0
                while True:
                    i = mm.find(pat, start)
                    if i < 0:
                        break
                    hit = i
                    break
                # proceed only if the constraint is loose? no - skip
                hit = None
            if hit is None:
                seen.add(key)
                print("no match for va=0x%x first8=%08x %08x" % (va, w0, w1))
                continue
            pa = (hit & ~0xFFF) + BASE
            seen.add(key)
            if inject(va & ~0xFFF, pa):
                injected += 1
        time.sleep(0.5)
    mm.close()
    f.close()
    print("done: injected=%d" % injected)

if __name__ == "__main__":
    main()