#!/usr/bin/env python3
"""Is SDMA's PTEPDE (page-table-entry write) privilege-gated like SRBM_WRITE?

`hw/sdma_srbm_probe.py` established [HW] that `SRBM_WRITE` runs only with
`RB_CNTL.RB_PRIV = 1`. PTEPDE is the other packet that matters for isolation: it writes
page-table entries, so a queue that can issue it can forge translations.

Static reading (`../AUDIT-SDMA-01.md`): the PTEPDE handler has **no privilege check**.
It does consult `reg[0x5a6c]` bits 0 and 8, and an earlier draft of that document called
this "PTEPDE's gate" -- that was wrong and is corrected there: the branch does not
reject anything, it selects between two *implementations* (the bits-set path applies an
extra address/value fixup via `_DMA_DATA_6`, then both paths fall into the same PTE
write at `_PKT_0xc00_0`).

So the prediction is: PTEPDE executes regardless of RB_PRIV. This tests it.

Packet (SDMA_PKT_PTEPDE_GEN, matching sdma_v6_0_vm_set_pte_pde):
    header, pe_lo, pe_hi, mask_lo, mask_hi, value_lo, value_hi, incr, 0, count-1
which is exactly the order the handler pulls off r1 at 0x09b4..0x09c2.

Method mirrors the SRBM probe: bring the queue up privileged, then clear *only*
RB_PRIV on the live queue, because a queue configured unprivileged from scratch does
not fetch at all and would make "nothing happened" meaningless.

Writes land in a VRAM buffer this script allocates. Nothing else is touched.

usage: PYTHONPATH=/root/psp/tinygrad sdma_ptepde_probe.py --apply
"""
import argparse
import struct
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from sdma_srbm_probe import Probe, SDMA_OP_NOP                    # noqa: E402

SDMA_OP_PTEPDE = 12
COUNT = 8


def ptepde_gen(pe_mc, value, incr, count):
    return [SDMA_OP_PTEPDE,
            pe_mc & 0xffffffff, (pe_mc >> 32) & 0xffffffff,
            0, 0,                                   # mask/flags
            value & 0xffffffff, (value >> 32) & 0xffffffff,
            incr, 0, count - 1]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--apply', action='store_true')
    args = ap.parse_args()
    if not args.apply:
        print('dry run: pass --apply')
        return 0

    from tinygrad import Device
    adev = Device["AMD"].iface.dev_impl
    p = Probe(adev)
    buf = adev.mm.palloc(0x1000, zero=True, boot=False)
    buf_mc = adev.paddr2mc(buf)
    print('target buffer paddr 0x%x mc 0x%x, %d entries of 8 bytes' % (buf, buf_mc, COUNT))

    results = {}
    for priv in (1, 0):
        print('\n=== RB_PRIV = %d' % priv)
        if priv == 1:
            p.setup(rb_priv=1)
        else:
            p.reg('_RB_CNTL').update(rb_priv=0)      # live flip, ring stays configured
        print('  RB_CNTL.rb_priv = %s' % p.reg('_RB_CNTL').read_bitfields().get('rb_priv'))

        value = 0xAAAA0000BBBB0000 if priv else 0xCCCC0000DDDD0000
        adev.vram.view(buf, 0x40)[:] = bytes(0x40)
        adev.gmc.flush_hdp()
        before = bytes(adev.vram.view(buf, 0x40)[:])

        ok, dt = p.submit([SDMA_OP_NOP])
        print('  control: NOP consumed=%s (%.1f ms)' % (ok, 1000 * dt))
        if not ok:
            print('  !! queue not running -- inconclusive for this setting')
            results[priv] = None
            continue

        ok2, dt2 = p.submit(ptepde_gen(buf_mc, value, 0x1000, COUNT))
        after = bytes(adev.vram.view(buf, 0x40)[:])
        wrote = after != before
        ents = struct.unpack('<8Q', after)
        print('  PTEPDE consumed=%s (%.1f ms) -> buffer %s' % (
            ok2, 1000 * dt2, 'WRITTEN' if wrote else 'unchanged'))
        for i, e in enumerate(ents[:4]):
            print('     entry[%d] = 0x%016x' % (i, e))
        results[priv] = wrote

    print('\n=== result')
    if results.get(1) is None or results.get(0) is None:
        print('  inconclusive: a control failed')
    elif results[1] and results[0]:
        print('  PTEPDE IS NOT PRIVILEGE-GATED: page-table entries were written with')
        print('  RB_PRIV=0, on the same queue where SRBM_WRITE is refused.')
        print('  => the two packets have different trust models in this firmware.')
    elif results[1] and not results[0]:
        print('  PTEPDE is gated after all -- the static reading (no privilege check in')
        print('  the handler) is wrong or the gate lives outside the F32 code path.')
    else:
        print('  PTEPDE did not write even with RB_PRIV=1: addressing/setup is wrong,')
        print('  no conclusion can be drawn.')
    return 0


if __name__ == '__main__':
    raise SystemExit(main())