#!/usr/bin/env python3
# walk.py - host-side replica of the service's svc_gmmu_walk + pd_discover,
# run against /dev/shm/vgpu-guest-ram. Discovers the page directory that
# reproduces the service's map[] sample pairs, then walks the ring/window
# VAs through it.
import mmap, struct, sys

F = "/dev/shm/vgpu-guest-ram"
SEMA_PA_MAX = 0x180000000

# samples from the r690 dmesg map[] lines
SAMPLES = [
    (0x100000000, 0x100326000),
    (0x100e2f000, 0x10096e000),
    (0x100e33000, 0x10096b000),
    (0x100e34000, 0x10096a000),
    (0x100e35000, 0x100969000),
    (0x100e37000, 0x102dcb000),
    (0x100e38000, 0x102dca000),
    (0x100e39000, 0x102dc9000),
]

TARGET_VAS = [0x44a0000, 0x4500000, 0x4560000, 0x4480000, 0x4485000,
              0x4ca0000, 0x4ca5000, 0x4d00000, 0x4b17000, 0x4020000,
              0x4100000, 0x4160000]

mm = None

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

def rd(pa, n=8):
    if pa + n > SEMA_PA_MAX:
        return None
    try:
        d = mm[off(pa):off(pa)+n]
        return d if len(d) == n else None
    except Exception:
        return None

def rd64(pa):
    d = rd(pa, 8)
    return None if d is None else int.from_bytes(d, "little")

def pde_aperture(ap):
    if ap == 1: return "FB"
    if ap in (2, 3): return "SYS"
    return None

def pte_aperture(ap):
    if ap == 0: return "FB"
    if ap == 1: return None
    return "SYS"

def pte_addr(e, fb):
    if fb:
        return ((e >> 8) & 0x1FFFFFF) << 12
    return ((e >> 8) & 0x3FFFFFFFFFFF) << 12

def pte_leaf(table, idx):
    e = rd64(table + 8 * idx)
    if e is None or not (e & 1):
        return None
    ap = pte_aperture((e >> 1) & 0x3)
    if ap is None:
        return None
    a = pte_addr(e, ap == "FB")
    if not fb_ok(a):
        return None
    return a, ap

def fb_ok(a):
    return a + 8 <= SEMA_PA_MAX

def walk(va):
    table = None
    # root_bits = ilog2(entries=4) = 2; shift = 49-2 = 47
    e = rd64(table) if False else None
    # level 0
    idx = (va >> 47) & 0x3
    # root table is the candidate itself (passed separately)
    return None

def walk_from(pdb, va, trace=False):
    table = pdb
    for level in range(3):
        if level == 0:
            bits, shift = 2, 47
        else:
            bits, shift = 9, 38 if level == 1 else 29
        idx = (va >> shift) & ((1 << bits) - 1)
        e = rd64(table + 8 * idx)
        if e is None:
            return None
        if trace:
            print("   L%d table=0x%x idx=%u e=0x%016x" % (level, table, idx, e))
        if e & 1:
            return None
        ap = pde_aperture((e >> 1) & 0x3)
        if ap is None:
            return None
        table = pte_addr(e, ap == "FB")
        if not fb_ok(table):
            return None
    idx = (va >> 21) & 0xff
    big = rd64(table + 16 * idx)
    small = rd64(table + 16 * idx + 8)
    if big is None or small is None:
        return None
    if trace:
        print("   PD0 table=0x%x idx=%u big=0x%016x small=0x%016x" %
              (table, idx, big, small))
    if not (small & 1):
        sap = pde_aperture((small >> 1) & 0x3)
        if sap is not None:
            t = pte_addr(small, sap == "FB")
            r = pte_leaf(t, (va >> 12) & 0x1ff)
            if r:
                return r[0] + (va & 0xfff), r[1]
    if big & 1:
        bap = pte_aperture((big >> 1) & 0x3)
        if bap is not None:
            a = pte_addr(big, bap == "FB") + (va & ((1 << 21) - 1))
            if fb_ok(a):
                return a, bap
    else:
        bap = pde_aperture((big >> 1) & 0x3)
        if bap is not None:
            t = pte_addr(big, bap == "FB")
            r = pte_leaf(t, (va >> 16) & 0x1f)
            if r:
                return r[0] + (va & 0xffff), r[1]
    return None

def main():
    global mm
    f = open(F, "rb")
    mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    size = len(mm)
    print("# scanning for directory candidates...")
    cands = []
    for fo in range(0, size, 4096):
        e = int.from_bytes(mm[fo:fo+8], "little")
        if e and not (e & 1):
            ap = (e >> 1) & 0x3
            if ap in (1, 2, 3):
                pa = fo if fo < 0x80000000 else fo + 0x80000000
                if pa < SEMA_PA_MAX:
                    cands.append(pa)
    print("# candidates: %d" % len(cands))
    best = []
    for c in cands:
        hits = 0
        for (sva, spa) in SAMPLES:
            r = walk_from(c, sva)
            if r and (r[0] & ~0xfff) == (spa & ~0xfff):
                hits += 1
        if hits >= 2:
            best.append((hits, c))
    best.sort(reverse=True)
    print("# candidates with >=2 hits: %d" % len(best))
    for hits, c in best[:10]:
        print("  pdb=0x%x hits=%d/8" % (c, hits))
    if not best:
        return 1
    pdb = best[0][1]
    print("== walking targets through pdb=0x%x:" % pdb)
    for va in TARGET_VAS:
        r = walk_from(pdb, va)
        if r:
            print("   va=0x%x -> pa=0x%x ap=%s" % (va, r[0], r[1]))
        else:
            print("   va=0x%x -> no translation" % va)
    mm.close(); f.close()
    return 0

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