#!/usr/bin/env python3
"""Does the PSP bootloader enforce anti-rollback on SYS_DRV?

Unblocked in pass 144: historical `psp_13_0_0_sos.bin` -- the exact file this part loads --
yields older signed SYS_DRV builds for this ASIC:

    running  49.0.54  (0x310036)  161,680 B
    older    49.0.47  (0x31002f)  157,584 B
    older    49.0.37  (0x310025)  157,584 B

The bootloader accepts component loads during driver init (that is how the running firmware
got there), so this substitutes an older signed build in that window and boots.

  * boots -> the bootloader accepts a signed but OLDER component: no anti-rollback, and any
    vulnerability fixed between that build and the current one is reachable again
  * refused -> anti-rollback (or a version/fuse gate) is enforced; a boundary that holds

Before substituting, the extracted blob is checked for the same `$PS1` container magic at
+16 as the running component, so a parse error cannot masquerade as a rejection.

argv[1]: path to sos bin | 'identity' (re-inject the running bytes, positive control)
"""
import ctypes, hashlib, sys
sys.path.insert(0, '/root/psp/hw')
import tinygrad.runtime.support.am.amdev as amdev
import tinygrad.runtime.autogen.am.am as am

def extract_sysdrv(path):
    blob = open(path, 'rb').read()
    hdr = am.struct_psp_firmware_header_v2_0.from_buffer_copy(blob)
    ch = hdr.header
    d = am.struct_psp_fw_bin_desc.from_address(ctypes.addressof(hdr.psp_fw_bin))
    if d.fw_type != am.PSP_FW_TYPE_PSP_SYS_DRV:
        raise RuntimeError('first descriptor is not SYS_DRV (%#x)' % d.fw_type)
    start = d.offset_bytes + ch.ucode_array_offset_bytes
    return blob[start:start + d.size_bytes], d.fw_version

target = sys.argv[1]
sub = None; subver = None
if target != 'identity':
    sub, subver = extract_sysdrv(target)
    if sub[16:20] != b'$PS1':
        print('EXTRACT BAD: container magic at +16 is %r, not $PS1 -> refusing to boot with it'
              % sub[16:20]); sys.exit(2)
    print('substitute: %s ver=%#x len=%d sha=%s magic_ok=True'
          % (target.split('/')[-1], subver, len(sub), hashlib.sha256(sub).hexdigest()[:16]),
          flush=True)
else:
    print('identity run: re-inject the running bytes (positive control)', flush=True)

orig = amdev.AMFirmware.__init__
def patched(self, adev):
    orig(self, adev)
    running = bytes(self.sos_fw[am.PSP_FW_TYPE_PSP_SYS_DRV])
    print('  running SYS_DRV len=%d sha=%s'
          % (len(running), hashlib.sha256(running).hexdigest()[:16]), flush=True)
    newb = running if sub is None else sub
    self.sos_fw[am.PSP_FW_TYPE_PSP_SYS_DRV] = memoryview(bytearray(newb))
    print('  INJECTED len=%d sha=%s (%s)'
          % (len(newb), hashlib.sha256(newb).hexdigest()[:16],
             'identity' if sub is None else 'older build %#x' % subver), flush=True)
amdev.AMFirmware.__init__ = patched

from tinygrad import Device
try:
    adev = Device["AMD"].iface.dev_impl
    alive = bool(adev.psp.is_sos_alive())
    print('BOOTRESULT ok=True sos_alive=%s' % alive, flush=True)
    print('VERDICT %s' % ('ACCEPTED - booted with this SYS_DRV' if alive
                          else 'booted but sOS not alive'), flush=True)
except Exception as e:
    print('BOOTRESULT ok=False err=%s %s' % (type(e).__name__, str(e)[:160]), flush=True)
    print('VERDICT REFUSED - the bootloader would not come up with this SYS_DRV', flush=True)