#!/usr/bin/env python3
"""Require RT_ICON and RT_GROUP_ICON inside a PE, not a sibling .ico file."""

import struct
import sys
import zipfile

RT_ICON = 3
RT_GROUP_ICON = 14


def pe_has_icon(data: bytes) -> None:
    if data[:2] != b"MZ":
        raise SystemExit("executable does not start with MZ")
    e_lfanew = struct.unpack_from("<I", data, 0x3C)[0]
    if data[e_lfanew : e_lfanew + 4] != b"PE\0\0":
        raise SystemExit("executable has no PE signature")
    coff = e_lfanew + 4
    nsections = struct.unpack_from("<H", data, coff + 2)[0]
    optsize = struct.unpack_from("<H", data, coff + 16)[0]
    opt = coff + 20
    magic = struct.unpack_from("<H", data, opt)[0]
    if magic == 0x20B:
        dd_off = opt + 112
    elif magic == 0x10B:
        dd_off = opt + 96
    else:
        raise SystemExit(f"unknown PE magic {magic:#x}")
    res_rva, res_size = struct.unpack_from("<II", data, dd_off + 8 * 2)
    if res_rva == 0 or res_size == 0:
        raise SystemExit("PE has no resource directory")
    sections = []
    sec = opt + optsize
    for index in range(nsections):
        off = sec + index * 40
        vsize, va, rawsize, rawptr = struct.unpack_from("<IIII", data, off + 8)
        sections.append((va, rawptr, max(rawsize, vsize)))

    def rva_to_off(rva: int) -> int:
        for va, rawptr, span in sections:
            if va <= rva < va + span:
                return rawptr + (rva - va)
        raise SystemExit(f"resource RVA {rva:#x} is outside the sections")

    off = rva_to_off(res_rva)
    nnames, nids = struct.unpack_from("<HH", data, off + 12)
    found = set()
    base = off + 16
    for index in range(nnames + nids):
        name_or_id, nextv = struct.unpack_from("<II", data, base + index * 8)
        if name_or_id & 0x80000000:
            continue
        if nextv & 0x80000000 and name_or_id in (RT_ICON, RT_GROUP_ICON):
            found.add(name_or_id)
    missing = {RT_ICON, RT_GROUP_ICON} - found
    if missing:
        raise SystemExit(f"PE is missing icon resource types {sorted(missing)}")


def main() -> None:
    path = sys.argv[1]
    if path.endswith(".zip"):
        with zipfile.ZipFile(path) as archive:
            names = archive.namelist()
            loose = [name for name in names if name.lower().endswith(".ico")]
            if loose:
                raise SystemExit(f"zip carries a loose icon file: {loose}")
            exe_name = "paperdoll-renderer.exe"
            if exe_name not in names:
                raise SystemExit("zip has no paperdoll-renderer.exe")
            pe_has_icon(archive.read(exe_name))
    else:
        with open(path, "rb") as handle:
            pe_has_icon(handle.read())
    print(f"embedded icon ok: {path}")


if __name__ == "__main__":
    main()