#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# memprobe.py - read/write another process's memory through /proc/<pid>/mem
# (root only, no gdb/ptrace attach needed).
#
#   scan PID [--min 0x204800000 --max 0x204c0000] [--limit 8]
#       find tracker objects: qwords inside [min,max) (the client's sem VAs),
#       then objects whose +0x10 points at them (sem objects), then objects
#       whose +0x9430 points at those (the libcuda tracker at p+0x9410).
#   dump PID 0xP
#       dump the tracker at 0xP (fields sem/mirror/completed, sem value,
#       progress object at p+0x48 and its indexed slots).
#   poke PID 0xADDR=0xVAL,...
#   watch PID 0xADDR,... [--ms 2] [--secs 20]
import sys, os, time, array, struct

WS = 8

def rd(fd, addr, n):
    try:
        return os.pread(fd, n, addr)
    except OSError:
        return None

def wr(fd, addr, data):
    return os.pwrite(fd, data, addr)

def r32(buf, off=0):
    return struct.unpack_from("<I", buf, off)[0]
def r64(buf, off=0):
    return struct.unpack_from("<Q", buf, off)[0]

def regions(pid, writable=True):
    out = []
    with open("/proc/%d/maps" % pid) as f:
        for line in f:
            p = line.split()
            try:
                lo, hi = [int(x, 16) for x in p[0].split("-")]
            except Exception:
                continue
            perm = p[1]
            if writable and "w" not in perm:
                continue
            out.append((lo, hi, perm, " ".join(p[5:])))
    return out

def scan_qwords(pid, fd, lo, hi, needle_lo=None, needle_hi=None, needle_set=None, chunk=1 << 20):
    hits = []
    for rlo, rhi, perm, name in regions(pid):
        a = max(rlo, 0)
        while a < min(rhi, 1 << 47):
            n = min(chunk, rhi - a)
            buf = rd(fd, a, n)
            if not buf:
                a += n
                continue
            buf = buf[:n - (n % WS)]
            arr = array.array("Q")
            arr.frombytes(buf)
            if needle_set is not None:
                for i, v in enumerate(arr):
                    if v in needle_set:
                        hits.append((a + i * WS, v))
            else:
                for i, v in enumerate(arr):
                    if needle_lo is not None and v < needle_lo:
                        continue
                    if needle_hi is not None and v >= needle_hi:
                        continue
                    hits.append((a + i * WS, v))
            a += n
    return hits

def sem_objs(pid, fd, lo, hi):
    # A real sem object stores its GPU VA three times, at so+0x08, so+0x10 and
    # so+0x18. Requiring a duplicated neighbour keeps the scan selective: a
    # bare qword in the window matches thousands of unrelated pointers once
    # the ring/pushbuffer windows are part of the range.
    first = scan_qwords(pid, fd, None, None, needle_lo=lo, needle_hi=hi)
    cand = {}
    for addr, val in first:
        cand.setdefault(val, set()).add(addr)
    out = {}
    for val, addrs in cand.items():
        for a in addrs:
            buf = rd(fd, a - 8, 0x20)
            if not buf or len(buf) < 0x20:
                continue
            q0, q1, q2, q3 = (r64(buf, i) for i in (0, 8, 0x10, 0x18))
            if q1 == val and (q0 == val or q2 == val) and \
               (q2 == val or q3 == val):
                out.setdefault(val, set()).add(a - 0x10)
    return out

def find_trackers(pid, fd, lo, hi, limit=8):
    sems = sem_objs(pid, fd, lo, hi)
    print("sem VAs found: %d" % len(sems))
    for va, sos in sorted(sems.items())[:8]:
        print("  va=0x%x -> %d sem-obj candidate(s)" % (va, len(sos)))
    needles = set()
    for sos in sems.values():
        needles |= sos
    if not needles:
        return []
    hits = scan_qwords(pid, fd, None, None, needle_set=needles)
    pset = set()
    for addr, val in hits:
        cand = addr - 0x9430
        if cand >= 0:
            pset.add(cand)
    print("tracker candidates: %d" % len(pset))
    res = []
    for p in sorted(pset):
        buf = rd(fd, p + 0x9410, 0x40)
        if not buf or len(buf) < 0x40:
            continue
        sub = r64(buf, 0x08); last = r64(buf, 0x10); cur = r64(buf, 0x18); sem = r64(buf, 0x20)
        if sem not in needles:
            continue
        mem = None
        sob = rd(fd, sem, 0x40)
        if sob:
            mem = r64(sob, 0x10)
        res.append((p, sub, last, cur, sem, mem))
        if len(res) >= limit:
            break
    return res

def dump_tracker(pid, fd, p):
    t = p + 0x9410
    buf = rd(fd, t, 0x40)
    if not buf:
        print("tracker unreadable at 0x%x" % t); return
    print("TRK=0x%x sub=0x%x last=0x%x cur=0x%x sem_obj=0x%x"
          % (t, r64(buf, 8), r64(buf, 0x10), r64(buf, 0x18), r64(buf, 0x20)))
    sem = r64(buf, 0x20)
    sob = rd(fd, sem, 0x40)
    if sob:
        print("  sem_obj[0:0x40]=%s" % sob.hex())
        mem = r64(sob, 0x10)
        mvb = rd(fd, mem & ~0xfff, 0x1000)
        off = mem & 0xfff
        val = r32(mvb, off) if mvb and off + 4 <= len(mvb) else -1
        print("  sem_mem=0x%x val@+0x%x=0x%x" % (mem, off, val))
    q = r64(rd(fd, p + 0x48, 8) or b"\0" * 8)
    print("  prog_obj=0x%x next=0x%x cblist=0x%x idx95bc=0x%x p34=0x%x"
          % (q, r64(rd(fd, p + 0x9470, 8) or b"\0" * 8),
             r64(rd(fd, p + 0x9458, 8) or b"\0" * 8),
             r32(rd(fd, p + 0x95bc, 4) or b"\0" * 4),
             r32(rd(fd, p + 0x34, 4) or b"\0" * 4)))
    if q:
        qb = rd(fd, q, 0x50)
        if qb:
            size = r32(qb, 8); idx = r32(qb, 0x48); cur = r32(qb, 0x4c); arr = r64(qb, 0x40)
            print("    size=%d idx=0x%x cursor=0x%x arr=0x%x" % (size, idx, cur, arr))
            for s in (idx, cur):
                slot = arr + (s % size if size else s) * 0x38
                sb = rd(fd, slot, 0x38)
                print("    slot[%d]@0x%x=%s" % (s, slot, sb.hex() if sb else "unreadable"))

def main():
    if len(sys.argv) < 3:
        print(__doc__); return 1
    mode = sys.argv[1]
    pid = int(sys.argv[2])
    fd = os.open("/proc/%d/mem" % pid, os.O_RDWR)
    try:
        if mode == "scan":
            lo, hi = 0x204800000, 0x204c0000
            limit = 8
            brief = False
            for i, a in enumerate(sys.argv[3:]):
                if a == "--min": lo = int(sys.argv[4 + i], 16)
                if a == "--max": hi = int(sys.argv[4 + i], 16)
                if a == "--limit": limit = int(sys.argv[4 + i])
                if a == "--brief": brief = True
            res = find_trackers(pid, fd, lo, hi, limit)
            for p, sub, last, cur, sem, mem in res:
                print("P=0x%x sub=0x%x last=0x%x cur=0x%x sem_obj=0x%x sem_mem=0x%x" %
                      (p, sub, last, cur, sem, mem or 0))
                if not brief:
                    dump_tracker(pid, fd, p)
        elif mode == "dump":
            dump_tracker(pid, fd, int(sys.argv[3], 16))
        elif mode == "hex":
            a = int(sys.argv[3], 16)
            n = int(sys.argv[4], 16) if len(sys.argv) > 4 else 0x100
            d = rd(fd, a, n) or b""
            for off in range(0, len(d), 16):
                c = d[off:off + 16]
                q = [int.from_bytes(c[i:i + 8], "little") for i in (0, 8)]
                print("0x%x: %s %s" % (a + off, c.hex(), " ".join("0x%x" % x for x in q)))
        elif mode == "ctx":
            ctx = int(sys.argv[3], 16)
            count = r32(rd(fd, ctx + 0x10, 4) or b"\0" * 4)
            ent = r64(rd(fd, ctx + 0x18, 8) or b"\0" * 8)
            print("waitobj ctx=0x%x count=%d entries=0x%x" % (ctx, count, ent))
            for i in range(min(count, 48)):
                e = ent + i * 0x28
                eb = rd(fd, e, 0x28)
                if not eb:
                    print("  e%d unreadable" % i); continue
                typ = r32(eb, 0); p = r64(eb, 8); v = r64(eb, 0x10)
                x18 = r64(eb, 0x18); x20 = r64(eb, 0x20)
                print("  e%d type=0x%x p=0x%x v=0x%x x18=0x%x x20=0x%x" % (i, typ, p, v, x18, x20))
                if typ == 1 and p:
                    dump_tracker(pid, fd, p)
        elif mode == "poke":
            for spec in sys.argv[3].split(","):
                a, v = spec.split("=")
                a = int(a, 16); v = int(v, 16)
                n = wr(fd, a, v.to_bytes(8, "little"))
                print("poke 0x%x <- 0x%x (%s)" % (a, v, "ok" if n == 8 else "FAIL"))
        elif mode == "watch":
            addrs = [int(x, 16) for x in sys.argv[3].split(",")]
            ms = 2; secs = 20
            for i in range(4, len(sys.argv)):
                if sys.argv[i] == "--ms" and i + 1 < len(sys.argv):
                    ms = int(sys.argv[i + 1])
                if sys.argv[i] == "--secs" and i + 1 < len(sys.argv):
                    secs = int(sys.argv[i + 1])
            last = {}
            t0 = time.time()
            while time.time() - t0 < secs:
                for a in addrs:
                    b = rd(fd, a, 8)
                    if not b:
                        continue
                    v = r64(b)
                    if last.get(a) != v:
                        print("%.3f 0x%x = 0x%x" % (time.time() - t0, a, v))
                        last[a] = v
                time.sleep(ms / 1000.0)
        else:
            print(__doc__); return 1
    finally:
        os.close(fd)
    return 0

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