"""Machine-check the sos_proc_switch msg-headroom defect directly from firmware bytes.

Claim under test: sos_proc_switch (SVC 0xF2) validates the caller-supplied `msg`
pointer with only 0x28 bytes of headroom on its third window, while the callee
(SYS_DRV's sysdrv_message_handler) reads AND writes 0x7C bytes at that pointer.

This script extracts BOTH numbers from the actual images -- the slack constants
from sos2.elf and the access size from sys_drv2.elf -- and reports the mismatch.
Nothing is hardcoded from the write-up.
"""
import struct
from capstone import *
from capstone.arm import *

BASE = '/Users/user/dev/amd-psp-re/psp_out/handover/elf/'

def seg(path):
    raw = open(path, 'rb').read()
    ph = struct.unpack('<I', raw[28:32])[0]
    o, v, pa, f = struct.unpack('<IIII', raw[ph+4:ph+20])
    return raw[o:o+f], v

# ---- 1. slack constants inside sos_proc_switch -----------------------------
body, v = seg(BASE + 'sos2.elf')
FUNC, END = 0x22b4, 0x22b4 + 456
md = Cs(CS_ARCH_ARM, CS_MODE_THUMB); md.detail = True
subs, adds, wins = [], [], []
for i in md.disasm(body[FUNC-v:END-v], FUNC):
    if i.mnemonic.startswith('sub') and len(i.operands) == 3 and i.operands[2].type == ARM_OP_IMM:
        subs.append((i.address, i.operands[2].imm))
    if i.mnemonic.startswith('add') and len(i.operands) == 3 and i.operands[2].type == ARM_OP_IMM:
        adds.append((i.address, i.operands[2].imm))
    if i.mnemonic.startswith('mov') and len(i.operands) == 2 and i.operands[1].type == ARM_OP_IMM:
        if i.operands[1].imm in (0x200000, 0x300000, 0x400000, 0xD00000, 0xE00000):
            wins.append((i.address, i.operands[1].imm))

print("=== sos_proc_switch @ %#x ===" % FUNC)
print("window constants  :", [(hex(a), hex(c)) for a, c in wins])
slack_4k  = [(a, c) for a, c in subs if c == 0x1000]
slack_x28 = [(a, c) for a, c in adds if c == 0x28]
print("msg headroom via 'sub #0x1000' (type5/type1) :", [hex(a) for a, _ in slack_4k])
print("msg headroom via 'add #0x28'   (third path)  :", [hex(a) for a, _ in slack_x28])

# ---- 2. access size inside SYS_DRV's message handler -----------------------
body2, v2 = seg(BASE + 'sys_drv2.elf')
H, HEND = 0xE087E4, 0xE087E4 + 0x90
md2 = Cs(CS_ARCH_ARM, CS_MODE_THUMB); md2.detail = True
sizes = []
for i in md2.disasm(body2[H-v2:HEND-v2], H):
    if i.mnemonic.startswith('mov') and len(i.operands) == 2 and i.operands[1].type == ARM_OP_IMM:
        if i.operands[1].imm == 0x7C: sizes.append(hex(i.address))
print("\n=== sysdrv_message_handler @ %#x ===" % H)
print("literal 0x7C (frame access size) at :", sizes)

# ---- 3. the mismatch -------------------------------------------------------
WIN_END, SLACK, ACCESS = 0xE00000, 0x28, 0x7C
maxmsg = WIN_END - SLACK
print("\n=== RESULT ===")
print("  validation permits msg <= %#x" % maxmsg)
print("  callee accesses [msg, msg+%#x) -> ends at %#x" % (ACCESS, maxmsg + ACCESS))
print("  bytes past the window end      : %d" % (maxmsg + ACCESS - WIN_END))
print("  msg values that PASS but overflow: %#x .. %#x (%d)"
      % (WIN_END - ACCESS + 1, maxmsg, maxmsg - (WIN_END - ACCESS + 1) + 1))
ok = bool(slack_4k) and bool(slack_x28) and bool(sizes) and (SLACK < ACCESS)
print("\n%s" % ("*** HEADROOM DEFECT CONFIRMED FROM FIRMWARE BYTES ***" if ok
                else "could not confirm -- check offsets"))