#!/usr/bin/env python3
"""Fuzz the PRE-AUTHENTICATION IFWI directory parser, in real firmware.

Why this target. `rom_parse_directory` (0xE16964) is reached from TSU mailbox
command 0x04 (`tsu_parse_rom_image`, 0xE1AFD8) with a host-supplied buffer --
host-reachable at runtime, no reboot, no flash write -- and it runs BEFORE any
signature check. The signature gate that downgraded the 4.1c disclosure
(FINDING-TSU-CHAIN-E2E-IN-EMULATOR.md) sits AFTER this parser, so bugs here are
not covered by it. Its wrapping bounds checks are already demonstrated to
dereference out of bounds (FINDING-IFWI-WRAP-CONFIRMED-IN-EMULATOR.md). Nobody
has fuzzed it: the tree has seven fuzz harnesses and none for this.

No hardware, no flash, no card -- Unicorn only.

Detectors
  * FAULT      -- the parser dereferenced something it should not have
  * REDZONE    -- it wrote outside the caller's context buffer
  * HANG       -- instruction budget exhausted (unbounded walk)
  * ACCEPT     -- a mutated image still returns 0; interesting because the
                  parser is the gate that decides whether the writer runs

Controls, checked every run; a failure makes the result NOT REPORTABLE:
  * the unmutated image must parse (return 0)
  * a corrupted IFWI magic must be rejected
  * parser block coverage must be non-trivial, so a run that never reaches the
    parser cannot report "no faults"

    python3 harness/fuzz_ifwi_parser.py [iters] [seed]
"""
import collections
import os
import random
import struct
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))

from pspemu import sysdrv_env                                   # noqa: E402
from unicorn import UC_HOOK_BLOCK                               # noqa: E402

RPD = 0xE16964
PARSER_LO, PARSER_HI = 0xE0F700, 0xE17000      # parser + its helpers
IMG_VA = 0x10000000
CTX_SZ = 0x100
RED = b"\x5A\xFE" * 8
FLASH = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                     "..", "..", "psp_out", "flash_BIOS2_KNOWNGOOD_20260917.bin")

# Structure-aware mutation sites, from the parser's own reads
#   +20 / +44  EFS slot-A / slot-B directory pointers  (fw_tail_scan)
#   +0x1000    $SGN table pointer                      (fw_tail_record_check)
#   dir +8     entry count                             (param_list_scan)
INTERESTING = [0x00000000, 0x00000001, 0x00000004, 0x0000FFFF, 0x0001FFFF,
               0x00200000, 0x00200001, 0x7FFFFFFF, 0x80000000, 0xFFFFF000,
               0xFFFFFFF0, 0xFFFFFFFC, 0xFFFFFFFE, 0xFFFFFFFF]


class Fuzzer:
    def __init__(self, seed=0):
        self.rng = random.Random(seed)
        self.img = open(os.path.normpath(FLASH), "rb").read()
        self.stats = collections.Counter()
        self.findings = []
        self.blocks = set()
        sd = self.sd = sysdrv_env.make()
        sysdrv_env.run_init(sd)
        sd.uc.mem_map(IMG_VA, (len(self.img) + 0xFFF) & ~0xFFF)
        self.ctx = sd.alloc(CTX_SZ + 0x20)
        self.ctxred = self.ctx + CTX_SZ
        sd.uc.hook_add(UC_HOOK_BLOCK, lambda u, a, s, x: self.blocks.add(a),
                       begin=PARSER_LO, end=PARSER_HI - 1)
        sd.uc.ctl_flush_tb()

    def parse(self, image):
        sd = self.sd
        sd.uc.mem_write(IMG_VA, image)
        sd.write(self.ctx, b"\x00" * CTX_SZ)
        sd.write(self.ctxred, RED)
        try:
            rv = sd.run_function(RPD, (self.ctx, IMG_VA, 0, 0), count=8_000_000)
            fault = None
        except Exception as e:
            rv, fault = None, (sd.faults[-1] if getattr(sd, "faults", None) else {})
        red_ok = sd.read(self.ctxred, 16) == RED
        return rv, fault, red_ok

    def mutate(self):
        img = bytearray(self.img)
        desc = []
        for _ in range(self.rng.randint(1, 3)):
            kind = self.rng.random()
            if kind < 0.45:                       # structural field
                off = self.rng.choice([20, 44, 0x1000,
                                       0x2000 + 8, 0x3000 + 8,
                                       0x20000 + 8, 0x1B000 + 8])
                val = self.rng.choice(INTERESTING)
            elif kind < 0.75:                     # inside a directory/record
                off = self.rng.choice([0x2000, 0x3000, 0x20000, 0x1B000]) + \
                    4 * self.rng.randrange(0, 0x40)
                val = self.rng.choice(INTERESTING)
            else:                                 # free dword anywhere early
                off = 4 * self.rng.randrange(0, 0x8000)
                val = self.rng.getrandbits(32)
            struct.pack_into("<I", img, off, val)
            desc.append("[%#x]=%#010x" % (off, val))
        return bytes(img), ",".join(desc)

    def run_one(self):
        image, desc = self.mutate()
        rv, fault, red_ok = self.parse(image)
        if not red_ok:
            self.stats["REDZONE"] += 1
            self.findings.append(("REDZONE", desc, "-", "-"))
        if fault is not None:
            kind = fault.get("kind", "FAULT")
            self.stats[kind] += 1
            self.findings.append((kind, desc, hex(fault.get("pc", 0)),
                                  hex(fault.get("addr", 0))))
        elif rv == 0:
            self.stats["ACCEPTED"] += 1
            self.findings.append(("ACCEPTED", desc, "-", "-"))
        else:
            self.stats["rejected %#x" % (rv & 0xFFFFFFFF)] += 1


def main(iters=2000, seed=1):
    f = Fuzzer(seed)
    print("IFWI pre-auth parser fuzz: rom_parse_directory @%#x" % RPD)

    rv, fault, red = f.parse(f.img)
    print("CONTROL unmutated image      : rv=%s fault=%s redzone_ok=%s"
          % (rv if rv is None else hex(rv), fault is not None, red))
    if rv != 0:
        print("\n  Control failed: the parser does not accept the known-good image.")
        print("RESULT NOT REPORTABLE")
        return 2
    bad = bytearray(f.img); bad[0:4] = b"\xde\xad\xbe\xef"
    rv2, _, _ = f.parse(bytes(bad))
    print("CONTROL corrupted magic      : rv=%s (must be non-zero)"
          % (rv2 if rv2 is None else hex(rv2)))
    if rv2 == 0:
        print("\n  Control failed: a corrupted image is accepted, so the parser is")
        print("  not actually deciding anything here.")
        print("RESULT NOT REPORTABLE")
        return 2

    base_cov = len(f.blocks)
    for _ in range(iters):
        f.run_one()
    print("\niterations=%d seed=%d" % (iters, seed))
    for k, v in f.stats.most_common():
        print("   %-26s %d" % (k, v))
    print("   parser blocks covered      %d (control alone reached %d)"
          % (len(f.blocks), base_cov))
    if len(f.blocks) < 20:
        print("\n  Parser coverage is trivial -- a clean run here would say nothing.")
        print("RESULT NOT REPORTABLE")
        return 2

    interesting = [x for x in f.findings if x[0] not in ("ACCEPTED",)]
    if interesting:
        print("\ndistinct fault sites:")
        seen = set()
        for kind, desc, pc, addr in interesting:
            if (kind, pc) in seen:
                continue
            seen.add((kind, pc))
            print("   %-16s pc=%-12s addr=%-12s  via %s" % (kind, pc, addr, desc[:60]))
    else:
        print("\nno faults or redzone violations")
    return 0


if __name__ == "__main__":
    a = sys.argv[1:]
    sys.exit(main(int(a[0]) if a else 2000, int(a[1]) if len(a) > 1 else 1))