#!/usr/bin/env python3
"""Saturated-baseline fuzz: make every limit permissive, then sweep one offset field.

Why this profile and not the two already tried. `ta_fieldfuzz.py` says it best:

    "The earlier profiles filled every field at once, which a first-stage validator will
     usually reject, so deeper code was never reached. This mutates ONE dword per invoke..."

Both extremes have been run — all fields filled, and one field mutated over zeros. Neither
reaches the shape that actually matters here. `fw/FINDING-GCP-FE09-CONFIGINIT-HOST-OFFSET-WRITE.md`
showed GCP validating a host offset against **another host field**, so the write is reached only
when a *limit* field is permissive AND an *offset* field is hostile, together. A single-dword
mutation over zeros cannot produce that: with the limit at zero, the comparison rejects.

So: hold every limit-like field at a permissive value, then sweep each field through extremes.
That is the profile that defeats a host-versus-host check, and it is the profile under which the
one live TA-side defect found in this work is reachable. Confirming the harness can re-find that
defect is the positive control, and it runs first.

Oracle: PSP fault dump C2PMSG 82..88, then ring death, then a write past the declared window,
then stall, then a status differing from the saturated baseline.

usage: AM_RESET=1 PYTHONPATH=/root/psp/tinygrad ta_saturated_fuzz.py --apply gcp
"""
import argparse
import glob
import struct
import time

DECL = 0x4000
TOTAL = 0x8000
PAT = 0xA5A5A5A5

# fields GCP and friends read as offsets or limits (fw/FINDING-GCP-FE09-...md)
FIELDS = [0x04, 0x08, 0x0C, 0x10, 0x14, 0x18, 0x1C, 0x20, 0x24, 0x28,
          0x414,
          0x810, 0x814, 0x818, 0x81C, 0x820, 0x824, 0x828, 0x82C, 0x830, 0x834, 0x838,
          0x20 + 1012, 0x20 + 2040, 0x20 + 2044, 0x20 + 2048, 0x20 + 2052,
          0x20 + 2056, 0x20 + 2060, 0x20 + 2064, 0x20 + 2072, 0x20 + 2088,
          0x20 + 2092, 0x20 + 2100]
SATURATE = 0xFFFFFFFB          # large but below the overflow-sentinel the code tests for

# Saturation has a blind spot: a field that an equality test gates on must hold its exact value,
# or the handler returns before the interesting code. GCP's cmdFE09 needs *(params+0x820) == 1 to
# reach ConfigureInitialize, and with that field saturated it takes the auth-state path instead
# and returns 0x262. The first run of this harness therefore could not reach the one live defect
# in fw/FINDING-GCP-FE09-CONFIGINIT-HOST-OFFSET-WRITE.md -- its positive control silently passed
# by never executing the target. PIN restores such selectors.
PIN = {'gcp': {0x820: 1}}

# Saturation also breaks a "x + 4 <= limit" check: with both sides at SATURATE, x+4 overflows
# past the limit and the guard rejects. Reaching such code needs LIMIT fields large and OFFSET
# fields small, which one constant cannot express. Roles are only knowable where the static work
# has been done; for GCP they come from fw/FINDING-GCP-FE09-CONFIGINIT-HOST-OFFSET-WRITE.md.
ROLES = {'gcp': {'limits': [0x414, 0x818],
                 'offsets': [0x828, 0x82C, 0x830, 0x834, 0x838,
                             0x20 + 2040, 0x20 + 2052, 0x20 + 2056, 0x20 + 2060,
                             0x20 + 2064, 0x20 + 2088, 0x20 + 2092, 0x20 + 2100]}}
SWEEP = [0xFFFFFFFF, 0x7FFFFFFF, 0x10000000, 0x3BE8, 0x1000, 1]

TAS = {
    # 0xFE09 FIRST: it is the positive control. ConfigureInitialize writes at
    # params + 1048 + *(params+0x828), so the saturated baseline alone should drive that write
    # out of the declared window. If the control does not fire, the harness proves nothing.
    # In the first run it never executed, because 0xFE02 hung before the sweep reached it.
    'gcp':   ('ta_22-0-0-77_7c426e7f', [0xFE09, 1, 2, 4, 5, 0xFE07, 0xFE08,
                                       0xFE0A, 0xFE0B, 0xFE0C, 0xFE0D, 0xFE0F, 0xFE02]),
    'opm':   ('ta_25-0-0-81_dc596ba6', list(range(0, 11))),
    'merit': ('ta_24-0-0-70_5cd4afda', [1, 2, 3, 4, 5, 6]),
    'auc':   ('ta_16-0-0-9_9e9044cf', [1, 2, 3, 4, 5]),
    'fp':    ('ta_19-0-0-8_7e8be282', [1, 3, 4, 5, 6]),
    'csu':   ('ta_42-0-0-6_1c874c3c', [1, 2, 3]),
}


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

    from tinygrad import Device
    import tinygrad.runtime.autogen.am.am as am

    adev = Device["AMD"].iface.dev_impl
    psp = adev.psp
    FAULT = [getattr(adev, 'mmMP0_SMN_C2PMSG_%d' % n) for n in range(82, 89)]

    def faults():
        try:
            return tuple(r.read() for r in FAULT)
        except Exception:
            return None

    def alive():
        try:
            return bool(psp.is_sos_alive())
        except Exception:
            return None

    def alloc(n):
        pa = adev.mm.palloc(n, align=0x1000, zero=True, boot=False)
        return pa, adev.paddr2mc(pa)

    def load(blob):
        ap_, amc = alloc(0x80000)
        cp, cmc = alloc(TOTAL)
        n = (len(blob) + 3) & ~3
        adev.vram.view(ap_, n)[:] = blob + b'\x00' * (n - len(blob))
        adev.gmc.flush_hdp()
        c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TA)
        t = c.cmd.cmd_load_ta
        t.app_phy_addr_lo, t.app_phy_addr_hi = amc & 0xFFFFFFFF, amc >> 32
        t.app_len = len(blob)
        t.cmd_buf_phy_addr_lo, t.cmd_buf_phy_addr_hi = cmc & 0xFFFFFFFF, cmc >> 32
        t.cmd_buf_len = DECL
        try:
            r = psp._ring_submit(c)
            return r.resp.status & 0xFFFFFFFF, r.resp.session_id, cp
        except Exception as e:
            try:
                return int(str(e).strip().split()[-1]) & 0xFFFFFFFF, None, cp
            except Exception:
                return -1, None, cp

    def inv(sid, cp, cmd, off=None, val=None, pin=None, roles=None):
        """Saturated baseline: every field at SATURATE, pinned selectors restored, then one
        field overridden."""
        buf = bytearray(struct.pack('<I', PAT) * (TOTAL // 4))
        struct.pack_into('<I', buf, 0, cmd)
        if roles:
            # limits permissive, offsets zero -- the only shape that passes "x + 4 <= limit"
            for o in FIELDS:
                struct.pack_into('<I', buf, o, 0)
            for o in roles['limits']:
                struct.pack_into('<I', buf, o, SATURATE)
        else:
            for o in FIELDS:
                struct.pack_into('<I', buf, o, SATURATE)
        for o, v in (pin or {}).items():
            struct.pack_into('<I', buf, o, v)
        if off is not None:
            struct.pack_into('<I', buf, off, val)
        adev.vram.view(cp, TOTAL)[:] = bytes(buf)
        adev.gmc.flush_hdp()
        c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD)
        iv = c.cmd.cmd_invoke_cmd
        iv.session_id, iv.ta_cmd_id = sid, cmd
        t0 = time.time()
        try:
            st = psp._ring_submit(c).resp.status & 0xFFFFFFFF
        except Exception:
            st = -1
        dt = time.time() - t0
        adev.gmc.flush_hdp()
        out = bytes(adev.vram.view(cp, TOTAL)[:])
        bey = [o for o in range(DECL, TOTAL, 4)
               if struct.unpack_from('<I', out, o)[0] != PAT]
        return st, dt, bey

    f0 = faults()
    print('baseline fault regs=%s sos_alive=%s' % ([hex(x) for x in f0] if f0 else None, alive()),
          flush=True)
    grand = 0
    for name in args.tas.split(','):
        stem, cmds = TAS[name]
        cands = glob.glob('/root/psp/contall/' + stem + '*')
        if not cands:
            print('%-6s binary not found' % name, flush=True)
            continue
        pin = PIN.get(name, {})
        roles = ROLES.get(name)
        st, sid, cp = load(open(cands[0], 'rb').read())
        if st != 0 or sid is None:
            print('%-6s LOAD %#x' % (name, st), flush=True)
            continue
        hard = []
        n = 0
        for cmd in cmds:
            base, _, _ = inv(sid, cp, cmd, pin=pin, roles=roles)
            stop = False
            for off in (roles['offsets'] if roles else FIELDS):
                for v in SWEEP:
                    if off in pin:
                        continue
                    s2, dt, bey = inv(sid, cp, cmd, off, v, pin=pin, roles=roles)
                    n += 1
                    f = faults()
                    if f is not None and f0 is not None and f != f0:
                        hard.append(('FAULT-DUMP', hex(cmd), hex(off), hex(v),
                                     [hex(x) for x in f])); stop = True
                    elif s2 == -1:
                        hard.append(('RING-DEAD', hex(cmd), hex(off), hex(v))); stop = True
                    elif bey:
                        hard.append(('OUT-OF-WINDOW', hex(cmd), hex(off), hex(v),
                                     [hex(o) for o in bey[:4]])); stop = True
                    if stop:
                        break
                if stop:
                    break
            if stop:
                break
        grand += n
        print('%-6s %4d cases  %s' % (name, n, hard[:2] if hard else 'clean'), flush=True)
        if hard and hard[0][0] in ('FAULT-DUMP', 'RING-DEAD'):
            print('  stopping: %s' % hard[0][0], flush=True)
            break
    print('\ntotal %d cases; sos_alive=%s fault regs=%s'
          % (grand, alive(), [hex(x) for x in (faults() or [])]), flush=True)
    return 0


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