#!/usr/bin/env python3
"""Carve PSP firmware containers out of an AMD **Windows** driver package.

Motivation: the project's firmware corpus (`psp_out/versions/`, `versions_lfw/`) comes
from linux-firmware. AMD's Windows Adrenalin packages ship their own PSP blobs inside
`amdkmdag.sys`, and those are a *different* release train, so they may carry builds the
Linux corpus never had.

Carve rules are NOT reinvented here — they are the ones this project already got wrong
twice and then got right (`FINDING-FIRMWARE-CORPUS-WAS-40-PERCENT-COMPLETE.md`):

  * the signed header starts **0x10 before** the `$PS1` magic;
  * body     = container + 0x100, length `0x100 + size_fw_signed`;
  * container total = `0x200 + size_fw_signed`;
  * **the magic alone is not enough** — 110 of 204 `$PS1` matches in the SPI image were
    bogus (the magic occurs inside compressed/signed data with junk size fields), so a
    carve must also validate the module header `00 02 00 00` plus a plausible load base.

Header layout is AMD's `PSP_FW_IMAGE_HEADER` as recorded in `spi_parse.py`
(0x14 SizeFWSigned, 0x60 ImageVersion, 0x68 FirmwareLoadAddr, 0x6C SizeImage,
0x7C FwType, 0x7D FwSubType).

Blobs are classified into three buckets and only the first is treated as a clean
"module build":
  module      - valid `00 02 00 00` + known load base  -> comparable with versions/
  headerless  - plausible container, body has no module header (sOS sub-images: KDB,
                SPL, drivers, TAs) -> still hashed, but flagged
  rejected    - failed size/sanity checks

usage: carve_driver.py <file-or-dir> [--out DIR] [--json]
"""
import argparse
import hashlib
import json
import os
import struct
import sys

MAGIC = b'$PS1'
# Load bases seen in this part's modules (FINDING-FIRMWARE-CORPUS...md)
KNOWN_BASES = {0x00e00000: 'sys_drv', 0x00d00000: 'asd', 0x00400000: 'ta',
               0x00000000: 'sos?', 0x000c0000: 'bl?'}


def ver_str(raw4):
    return '.'.join(str(b) for b in raw4[::-1])


TA_VER_KEY = b'amd.ta.version\x00'
TA_DESC_KEY = b'gpd.ta.description\x00'
TA_APPID_KEY = b'gpd.ta.appID\x00'


def _after_key(body, key, want_printable=False, n=4):
    """Read the value that follows a property key in the TA property table.

    Verified against the corpus: hdcp_23.0.0.82 -> 52 00 00 17 -> 23.0.0.82, and
    asd_33.0.0.154 -> 9a 00 00 21 -> 33.0.0.154.  So the version is stored
    little-endian as {build, patch, minor, major}.
    """
    i = body.find(key)
    if i < 0:
        return None
    j = i + len(key)
    if not want_printable:
        v = body[j:j + n]
        return '%d.%d.%d.%d' % (v[3], v[2], v[1], v[0]) if len(v) == 4 else None
    import re as _re
    m = _re.search(rb'[ -~]{4,64}', body[j:j + 0x120])
    return m.group(0).decode() if m else None


def identify(body):
    """TA-level identity: the version AMD names builds by, plus a readable name."""
    return dict(ta_version=_after_key(body, TA_VER_KEY),
                description=_after_key(body, TA_DESC_KEY, want_printable=True),
                app_id=_after_key(body, TA_APPID_KEY, want_printable=True))


def parse_header(blob):
    if len(blob) < 0x100:
        return None
    g32 = lambda o: struct.unpack_from('<I', blob, o)[0]
    return dict(size_fw_signed=g32(0x14), sec_patch_level=g32(0x4C),
                image_version=blob[0x60:0x64], fw_load_addr=g32(0x68),
                size_image=g32(0x6C), fw_type=blob[0x7C], fw_subtype=blob[0x7D])


def carve(data, src):
    """Yield dicts for every plausible container in `data`."""
    out = []
    pos = 0
    while True:
        m = data.find(MAGIC, pos)
        if m < 0:
            break
        pos = m + 4
        cont = m - 0x10                      # signed header starts 0x10 before the magic
        if cont < 0:
            continue
        hdr = parse_header(data[cont:cont + 0x100])
        if not hdr:
            continue
        n = hdr['size_fw_signed']
        # size sanity: must be non-zero and fit (the bogus-match filter)
        if not (0 < n < 0x400000) or cont + 0x200 + n > len(data):
            out.append(dict(src=src, off=cont, verdict='rejected',
                            reason='size %#x' % n))
            continue
        body = data[cont + 0x100: cont + 0x100 + 0x100 + n]
        mod = body[:4] == b'\x00\x02\x00\x00'
        base = struct.unpack_from('<I', body, 4)[0] if len(body) >= 8 else None
        rec = dict(src=src, off=cont, size=len(body),
                   version=ver_str(hdr['image_version']),
                   fw_type=hdr['fw_type'], fw_subtype=hdr['fw_subtype'],
                   load_addr=hdr['fw_load_addr'],
                   patch=hdr['sec_patch_level'],
                   sha256=hashlib.sha256(body).hexdigest())
        if mod and base in KNOWN_BASES:
            rec.update(verdict='module', guess=KNOWN_BASES[base], body_base=base)
        elif mod:
            rec.update(verdict='module', guess='base_%#x' % base, body_base=base)
        else:
            rec.update(verdict='headerless', guess=None, body_base=None)
        rec.update(identify(body))
        rec['_body'] = body
        # the FULL signed container (header + body + trailer) -- this is what
        # LOAD_TA must be given, since the PSP authenticates the signature.
        rec['_container'] = data[cont:cont + 0x200 + n]
        out.append(rec)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('target')
    ap.add_argument('--out', help='write carved bodies here')
    ap.add_argument('--json', action='store_true')
    ap.add_argument('--min-size', type=int, default=4096)
    args = ap.parse_args()

    files = []
    if os.path.isdir(args.target):
        for root, _, fs in os.walk(args.target):
            for f in fs:
                p = os.path.join(root, f)
                try:
                    if os.path.getsize(p) >= args.min_size:
                        files.append(p)
                except OSError:
                    pass
    else:
        files = [args.target]

    recs = []
    for p in files:
        try:
            with open(p, 'rb') as fh:
                data = fh.read()
        except Exception:
            continue
        if MAGIC not in data:
            continue
        recs.extend(carve(data, p))

    mods = [r for r in recs if r['verdict'] == 'module']
    hl = [r for r in recs if r['verdict'] == 'headerless']
    rej = [r for r in recs if r['verdict'] == 'rejected']
    print('scanned %d files: %d module, %d headerless, %d rejected'
          % (len(files), len(mods), len(hl), len(rej)))

    seen = {}
    for r in sorted(mods, key=lambda r: (r['guess'] or '', r['version'])):
        k = r['sha256']
        if k in seen:
            continue
        seen[k] = r
        print('  %-9s hdr=%-11s ta=%-13s %-26s %#08x %s'
              % (r['guess'], r['version'], r['ta_version'] or '-',
                 (r['description'] or '-')[:26], r['size'], r['sha256'][:10]))
    print('unique module bodies: %d' % len(seen))

    if args.out:
        os.makedirs(args.out, exist_ok=True)
        for r in seen.values():
            name = '%s_%s_%s.bin' % (r['guess'], r['version'], r['sha256'][:8])
            with open(os.path.join(args.out, name), 'wb') as fh:
                fh.write(r['_body'])
        # headerless bodies too, kept separate and clearly labelled
        hd = os.path.join(args.out, 'headerless')
        os.makedirs(hd, exist_ok=True)
        for r in {x['sha256']: x for x in hl}.values():
            with open(os.path.join(hd, 'hl_%s_%s.bin'
                                   % (r['version'], r['sha256'][:8])), 'wb') as fh:
                fh.write(r['_body'])
        print('wrote %d module + %d headerless bodies to %s'
              % (len(seen), len({x['sha256'] for x in hl}), args.out))

    if args.json:
        for r in recs:
            r.pop('_body', None)
        print(json.dumps(recs, indent=1))
    return 0


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