#!/usr/bin/env python3
"""Extract Windows ICO resources from a PE executable.

This uses only the Python standard library. It reads RT_GROUP_ICON and
RT_ICON resources, then rebuilds a normal .ico file.
"""
from __future__ import annotations

import argparse
import struct
from pathlib import Path


RT_ICON = 3
RT_GROUP_ICON = 14


class PeIconError(RuntimeError):
    pass


def u16(data: bytes, offset: int) -> int:
    return struct.unpack_from("<H", data, offset)[0]


def u32(data: bytes, offset: int) -> int:
    return struct.unpack_from("<I", data, offset)[0]


def parse_sections(data: bytes) -> tuple[int, list[dict[str, int]]]:
    if data[:2] != b"MZ":
        raise PeIconError("not an MZ executable")
    pe_offset = u32(data, 0x3C)
    if data[pe_offset:pe_offset + 4] != b"PE\0\0":
        raise PeIconError("missing PE signature")
    file_header = pe_offset + 4
    section_count = u16(data, file_header + 2)
    optional_header_size = u16(data, file_header + 16)
    optional = file_header + 20
    magic = u16(data, optional)
    if magic == 0x10B:
        data_directory = optional + 96
    elif magic == 0x20B:
        data_directory = optional + 112
    else:
        raise PeIconError(f"unsupported optional header magic: 0x{magic:04x}")
    resource_rva = u32(data, data_directory + 2 * 8)
    sections = []
    section_table = optional + optional_header_size
    for index in range(section_count):
      row = section_table + index * 40
      sections.append({
          "virtual_size": u32(data, row + 8),
          "virtual_address": u32(data, row + 12),
          "raw_size": u32(data, row + 16),
          "raw_pointer": u32(data, row + 20),
      })
    return resource_rva, sections


def rva_to_offset(rva: int, sections: list[dict[str, int]]) -> int:
    for section in sections:
        start = section["virtual_address"]
        size = max(section["virtual_size"], section["raw_size"])
        if start <= rva < start + size:
            return section["raw_pointer"] + (rva - start)
    raise PeIconError(f"RVA not mapped to a section: 0x{rva:08x}")


def resource_entries(data: bytes, directory_offset: int) -> list[tuple[int | str, int, bool]]:
    named_count = u16(data, directory_offset + 12)
    id_count = u16(data, directory_offset + 14)
    entries = []
    cursor = directory_offset + 16
    for _index in range(named_count + id_count):
        name_value = u32(data, cursor)
        data_value = u32(data, cursor + 4)
        cursor += 8
        if name_value & 0x80000000:
            name = f"name@0x{name_value & 0x7fffffff:08x}"
        else:
            name = name_value & 0xffff
        entries.append((name, data_value & 0x7fffffff, bool(data_value & 0x80000000)))
    return entries


def collect_resources(data: bytes, resource_base: int, sections: list[dict[str, int]]) -> dict[int, dict[int | str, list[dict[str, bytes | int]]]]:
    resources: dict[int, dict[int | str, list[dict[str, bytes | int]]]] = {}
    root = resource_entries(data, resource_base)
    for type_id, type_offset, type_is_directory in root:
        if not type_is_directory or not isinstance(type_id, int):
            continue
        type_dir = resource_base + type_offset
        for name_id, name_offset, name_is_directory in resource_entries(data, type_dir):
            if not name_is_directory:
                continue
            name_dir = resource_base + name_offset
            for lang_id, data_offset, lang_is_directory in resource_entries(data, name_dir):
                if lang_is_directory:
                    continue
                entry = resource_base + data_offset
                rva = u32(data, entry)
                size = u32(data, entry + 4)
                file_offset = rva_to_offset(rva, sections)
                payload = data[file_offset:file_offset + size]
                resources.setdefault(type_id, {}).setdefault(name_id, []).append({
                    "lang": int(lang_id) if isinstance(lang_id, int) else -1,
                    "rva": rva,
                    "size": size,
                    "data": payload,
                })
    return resources


def build_ico(group_data: bytes, icons: dict[int | str, list[dict[str, bytes | int]]]) -> bytes:
    if len(group_data) < 6:
        raise PeIconError("group icon resource too small")
    reserved, icon_type, count = struct.unpack_from("<HHH", group_data, 0)
    if reserved != 0 or icon_type != 1:
        raise PeIconError("resource is not an icon group")
    entries = []
    payloads = []
    cursor = 6
    for _index in range(count):
        if cursor + 14 > len(group_data):
            raise PeIconError("truncated group icon entry")
        width, height, color_count, reserved_byte, planes, bit_count, size, icon_id = struct.unpack_from(
            "<BBBBHHIH",
            group_data,
            cursor,
        )
        cursor += 14
        icon_rows = icons.get(icon_id)
        if not icon_rows:
            raise PeIconError(f"missing RT_ICON id {icon_id}")
        payload = icon_rows[0]["data"]
        if not isinstance(payload, bytes):
            raise PeIconError(f"invalid RT_ICON payload id {icon_id}")
        entries.append((width, height, color_count, reserved_byte, planes, bit_count, len(payload)))
        payloads.append(payload)
    header = bytearray(struct.pack("<HHH", 0, 1, len(entries)))
    offset = 6 + len(entries) * 16
    for width, height, color_count, reserved_byte, planes, bit_count, size in entries:
        header.extend(struct.pack("<BBBBHHII", width, height, color_count, reserved_byte, planes, bit_count, size, offset))
        offset += size
    return bytes(header) + b"".join(payloads)


def score_group(group_data: bytes) -> tuple[int, int]:
    if len(group_data) < 6:
        return (0, 0)
    count = u16(group_data, 4)
    max_side = 0
    cursor = 6
    for _index in range(count):
        if cursor + 14 > len(group_data):
            break
        width = group_data[cursor] or 256
        height = group_data[cursor + 1] or 256
        max_side = max(max_side, width, height)
        cursor += 14
    return (max_side, count)


def extract_icon(exe: Path, out: Path, group_id: int | None = None) -> dict[str, int | str]:
    data = exe.read_bytes()
    resource_rva, sections = parse_sections(data)
    resource_base = rva_to_offset(resource_rva, sections)
    resources = collect_resources(data, resource_base, sections)
    groups = resources.get(RT_GROUP_ICON, {})
    icons = resources.get(RT_ICON, {})
    if not groups or not icons:
        raise PeIconError("no icon resources found")
    if group_id is not None:
        group_rows = groups.get(group_id)
        if not group_rows:
            raise PeIconError(f"group icon id {group_id} not found")
        selected_id = group_id
        selected_group = group_rows[0]["data"]
    else:
        selected_id, selected_rows = max(
            groups.items(),
            key=lambda item: score_group(item[1][0]["data"] if isinstance(item[1][0]["data"], bytes) else b""),
        )
        selected_group = selected_rows[0]["data"]
    if not isinstance(selected_group, bytes):
        raise PeIconError("invalid group icon payload")
    ico = build_ico(selected_group, icons)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_bytes(ico)
    return {
        "groupId": int(selected_id) if isinstance(selected_id, int) else str(selected_id),
        "groupCount": len(groups),
        "iconCount": len(icons),
        "bytes": len(ico),
        "out": str(out),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("exe", type=Path)
    parser.add_argument("out", type=Path)
    parser.add_argument("--group-id", type=int)
    args = parser.parse_args()
    result = extract_icon(args.exe, args.out, group_id=args.group_id)
    print(
        f"wrote {result['out']} · group={result['groupId']} "
        f"groups={result['groupCount']} icons={result['iconCount']} bytes={result['bytes']}"
    )


if __name__ == "__main__":
    main()