"""Parse the SLWR register allowlist -- the firmware's own list of the SMN
registers a TA (and therefore the host, via RAS) is permitted to touch.

Why this matters: SYS_DRV gates every TA register access through
`reg_allowlist_lookup` (0xE17EA4), which consults a hash table at 0xE6D010.
That table is in BSS -- it is EMPTY in the image and populated at runtime -- so
the allowlist could not be read out of sys_drv2.elf. It is a separate signed
blob: `psp_out/drivers/rl.bin`, whose body begins with the magic 'SLWR'.

Format (recovered here):
    0x000  PSP container header (0x100 bytes; '$PS1' at +0x10, size at +0x14)
    0x110  'SLWR'
    0x114  u32 entry count
    0x120  entry[] : { u32 smn_address ; u32 granted_attrs }

Per STALEUS (see ARCH-INSIGHTS-FROM-STALEUS-PAPER.md) an SMN address splits as
upper 12 bits = IP selector, lower 20 bits = offset within that IP.

Run: python3 slwr_parse.py [path-to-rl.bin]
"""
import struct
import sys
from collections import Counter, defaultdict

PATH = sys.argv[1] if len(sys.argv) > 1 else "drivers/rl.bin"


def parse(path=PATH):
    d = open(path, "rb").read()
    base = d.find(b"SLWR")
    if base < 0:
        raise SystemExit("no SLWR magic in %s" % path)
    n = struct.unpack_from("<I", d, base + 4)[0]
    tbl = base + 0x10
    if tbl + 8 * n > len(d):
        raise SystemExit("entry count %d overruns the file" % n)
    ents = [struct.unpack_from("<II", d, tbl + 8 * i) for i in range(n)]
    return d, base, n, ents


def main():
    d, base, n, ents = parse()
    print("file        : %s (%#x bytes)" % (PATH, len(d)))
    print("SLWR magic  : %#x   entries: %d (%#x)" % (base, n, n))
    print("table       : %#x .. %#x" % (base + 0x10, base + 0x10 + 8 * n))
    print("addr range  : %#010x .. %#010x" % (min(a for a, _ in ents),
                                              max(a for a, _ in ents)))
    attrs = Counter(x for _, x in ents)
    print("attributes  : %s" % {hex(k): v for k, v in attrs.most_common()})
    if len(attrs) == 1:
        a = next(iter(attrs))
        print("              -> UNIFORM: every allowlisted register grants the")
        print("                 same attrs (byte0=%#04x, byte1=%#04x), so"
              % (a & 0xFF, (a >> 8) & 0xFF))
        print("                 attr_permits_usage applies one policy to all.")
    print()

    byip = defaultdict(list)
    for a, x in ents:
        byip[a >> 20].append(a & 0xFFFFF)
    print("%-8s %-6s %s" % ("IP", "count", "offsets"))
    for ip in sorted(byip):
        offs = sorted(set(byip[ip]))
        s = " ".join("%#07x" % o for o in offs[:6])
        print("%-8s %-6d %s%s" % (hex(ip), len(byip[ip]), s,
                                  " ... (%d distinct)" % len(offs) if len(offs) > 6 else ""))
    print()

    # structural decomposition of the large instances
    big = [ip for ip in byip if len(byip[ip]) > 4]
    if big:
        offs = sorted(set(byip[big[0]]))
        # find the two replication strides
        roots = sorted({o for o in offs if o < 0x1000})
        print("large-instance decomposition (IP %s):" % hex(big[0]))
        print("   %d root registers: %s" % (len(roots), [hex(o) for o in roots]))
        for stride in (0x1000, 0x10000):
            ok = all((o + stride) in offs for o in roots)
            print("   replicated at +%#x : %s" % (stride, ok))
        print("   => %d roots x 2 x 2 x %d instances = %d, plus pairs = %d"
              % (len(roots), len(big), len(roots) * 4 * len(big), n))
        print("   (a channel / sub-channel layout, i.e. memory-controller shaped)")
    print()
    print("NOTE these are addresses in the PSP's SMN numbering. They are NOT the")
    print("bases our host-side probes used (umc_survey.py takes those from the IP")
    print("discovery table via adev.regs_offset). See")
    print("FINDING-SLWR-ALLOWLIST-EXTRACTED.md.")


if __name__ == "__main__":
    main()