#!/usr/bin/env python3
# pt_bind.py - bind GMMU PT pages to VA bases and read out VA->PA for the
# client's ring/window family. Strategy:
#   1. ring-page candidates: pages whose 8B-aligned u64 pairs look like
#      GPFIFO entries (e0 = push VA & ~3 in the 0x44x-0x4ex families with
#      va<4G; e1 top byte 0; length bits sane).
#   2. anchor: for each candidate PA and each family VA, check every 4KB
#      page for entry idx=(VA>>12)&0x1ff whose PA part == candidate PA.
#      A hit binds (PT page, VA, PA).
#   3. read out: from bound PT pages, decode entries for the whole VA list.
import mmap, struct, sys

F = "/dev/shm/vgpu-guest-ram"
FAMILY = [
    0x4020000, 0x4100000, 0x4160000,           # r689 eng rings
    0x44a0000, 0x4500000, 0x4560000,           # r690 ring VAs
    0x4480000, 0x4485000, 0x4480980, 0x4480f1c,
    0x4ca0000, 0x4ca5000, 0x4d00000, 0x4b17000,
]

def gpa_to_off(pa):
    return pa if pa < 0x80000000 else pa - 0x80000000

def off_to_gpa(off):
    return off if off < 0x80000000 else off + 0x80000000

def pte_pa(v):
    return (v >> 8) << 12

def main():
    f = open(F, "rb")
    mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    size = len(mm)
    print("# file size 0x%x" % size)

    # ---- pass 1: ring-page candidates by content ----
    cands = []
    pos = 0
    step = 4096
    while pos + step <= size:
        d = mm[pos:pos+step]
        hits = []
        for i in range(0, step, 8):
            e0 = int.from_bytes(d[i:i+4], "little")
            e1 = int.from_bytes(d[i+4:i+8], "little")
            if 0x4400000 <= e0 <= 0x4e00000 and (e0 & 3) == 0 and (e1 & 0xff) == 0:
                lw = (e1 >> 10) & 0x1fffff
                if 0x40 <= lw <= 0x100000:
                    hits.append((i, e0, lw << 2))
        if len(hits) >= 3:
            cands.append((off_to_gpa(pos), hits))
            if len(cands) >= 12:
                break
        pos += step
    print("# ring-page candidates: %d" % len(cands))
    for pa, hits in cands:
        print("cand pa=0x%x hits=%d first=%s" %
              (pa, len(hits), ["0x%x/0x%x" % (h[1], h[2]) for h in hits[:4]]))
    if not cands:
        return 1

    # ---- pass 2: anchor PT pages for candidate PAs x family VAs ----
    idxs = {}
    for va in FAMILY:
        idxs.setdefault((va >> 12) & 0x1ff, []).append(va)
    cand_pas = set(pa for pa, _ in cands)
    bound = []
    pos = 0
    while pos + step <= size:
        pg = off_to_gpa(pos)
        # skip scanning the candidate pages themselves
        if pg not in cand_pas:
            d = mm[pos:pos+step]
            for idx, vas in idxs.items():
                o = idx * 8
                v = int.from_bytes(d[o:o+8], "little")
                if v & 1:
                    pa = pte_pa(v)
                    if pa in cand_pas:
                        bound.append((pg, idx, vas, pa, v & 0xff))
        pos += step
    print("# PT anchors found: %d" % len(bound))
    for pg, idx, vas, pa, fl in bound[:20]:
        print("PT pa=0x%x entry[0x%x] flags=0x%x -> PA=0x%x  (VA cands %s)" %
              (pg, idx, fl, pa, ["0x%x" % x for x in vas]))

    # ---- pass 3: read out PT entries for the whole family from each bound PT page ----
    seen_pt = set()
    for pg, idx, vas, pa, fl in bound:
        if pg in seen_pt:
            continue
        seen_pt.add(pg)
        off = gpa_to_off(pg)
        d = mm[off:off+4096]
        print("== PT page pa=0x%x family readout:" % pg)
        for va in FAMILY:
            i = (va >> 12) & 0x1ff
            v = int.from_bytes(d[i*8:i*8+8], "little")
            if v & 1:
                print("   va=0x%x idx=0x%x -> pa=0x%x (raw=0x%x fl=0x%x)" %
                      (va, i, pte_pa(v), v, v & 0xff))
            else:
                print("   va=0x%x idx=0x%x -> absent (raw=0x%x)" % (va, i, v))
    mm.close(); f.close()
    return 0

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