#!/usr/bin/env python3
# pdbtick.py - content-locate the GR0 ring page (deterministic first three
# GPFIFO entries) in the guest RAM file, then hand the verified
# {ring VA -> ring PA} pair to the service (sema_pdb_discover), which
# verifies the parked COPY_SERVER_PDES candidates and falls back to the
# vidmem walk-root sweep. Runs while the client is alive.
import mmap, struct, subprocess, sys, threading, time

F = "/dev/shm/vgpu-guest-ram"
PARAM = "/sys/module/vgsp/parameters/sema_pdb_discover"
PAT = struct.pack("<QQQ", 0x0005980004ca0000, 0x8000040004ca0598,
                  0x0003d00004ca059c)
RING_VA = 0x4cc0000

def gpa_of(x):
    return x if x < 0x80000000 else x - 0x80000000 + 0x100000000

def main():
    fobj = open(F, "rb")
    mm = mmap.mmap(fobj.fileno(), 0, access=mmap.ACCESS_READ)
    size = len(mm)
    nthr = 8
    step = size // nthr
    res = []
    lock = threading.Lock()
    def worker(a, b):
        CH = 8 << 20
        off = a
        while off < b:
            d = mm[off:off + CH]
            i = d.find(PAT)
            while i >= 0:
                with lock:
                    res.append(off + i)
                i = d.find(PAT, i + 1)
            off += CH - 64
    ts = []
    for k in range(nthr):
        a = k * step
        b = (k + 1) * step if k < nthr - 1 else size
        t = threading.Thread(target=worker, args=(a, b)); t.start(); ts.append(t)
    for t in ts: t.join()
    if not res:
        print("RING NOT FOUND")
        return None
    hits = [gpa_of(x) for x in res]
    print("ring candidates:", [hex(h) for h in hits], flush=True)
    for gpa in hits:
        cmd = "echo 0x%x,0x%x > %s" % (RING_VA, gpa, PARAM)
        r = subprocess.run(["sudo", "sh", "-c", cmd], capture_output=True)
        print("sema_pdb_discover 0x%x,0x%x rc=%d" % (RING_VA, gpa, r.returncode), flush=True)
    return hits


def loop():
    import time as _t
    dur = float(sys.argv[1]) if len(sys.argv) > 1 else 120
    end = _t.time() + dur
    fired = set()
    while _t.time() < end:
        hits = main()
        if hits:
            new = [h for h in hits if h not in fired]
            if new:
                fired.update(new)
        _t.sleep(1.0)
    print("# pdbtick end", flush=True)

loop()