"""PE icon guard: Windows PE32/PE32+ resources, and missing/broken groups."""
import pathlib
import shutil
import struct
import subprocess
import tempfile
import unittest

SCRIPT = pathlib.Path(__file__).resolve().parents[1] / "windows/check-icon.ps1"


def executable(pe64=True, icon=True, broken=False):
    data = bytearray(1024)
    def u16(at, value): struct.pack_into("<H", data, at, value)
    def u32(at, value): struct.pack_into("<I", data, at, value)
    u16(0, 0x5A4D)
    u32(0x3C, 0x80)
    u32(0x80, 0x4550)
    u16(0x86, 1)
    optional = 0x98
    size = 240 if pe64 else 224
    u16(0x94, size)
    u16(optional, 0x20B if pe64 else 0x10B)
    directories = optional + (112 if pe64 else 96)
    u32(directories + 16, 0x1000)
    u32(directories + 20, 512)
    section = optional + size
    u32(section + 12, 0x1000)
    u32(section + 16, 512)
    u32(section + 20, 512)
    base = 512
    def directory(relative, entries):
        u16(base + relative + 14, len(entries))
        for i, (name, offset) in enumerate(entries):
            u32(base + relative + 16 + 8*i, name)
            u32(base + relative + 20 + 8*i, offset)
    if icon:
        directory(0, [(3, 0x80000030), (14, 0x80000070)])
        directory(0x30, [(1, 0x80000050)])
        directory(0x50, [(1033, 0xB0)])
        directory(0x70, [(1, 0x80000090)])
        directory(0x90, [(1033, 0xC0)])
        u32(base + 0xB0, 0x1100)
        u32(base + 0xB4, 40)
        u32(base + 0xC0, 0x1140)
        u32(base + 0xC4, 20)
        u16(base + 0x142, 1)
        u16(base + 0x144, 1)
        u32(base + 0x14E, 40)
        u16(base + 0x152, 2 if broken else 1)
    return data


class IconGuard(unittest.TestCase):
    def run_guard(self, data):
        with tempfile.TemporaryDirectory() as folder:
            path = pathlib.Path(folder) / "app.exe"
            path.write_bytes(data)
            return subprocess.run([shutil.which("pwsh") or "powershell", "-NoProfile", "-File", str(SCRIPT), "-Path", str(path)], capture_output=True, text=True)

    def test_valid_pe32_and_pe64(self):
        for pe64 in (False, True):
            result = self.run_guard(executable(pe64))
            self.assertEqual(result.returncode, 0, result.stderr)
            self.assertIn("1 images, 1 group", result.stdout)

    def test_generic_executable_is_rejected(self):
        result = self.run_guard(executable(icon=False))
        self.assertNotEqual(result.returncode, 0)
        self.assertIn("RT_ICON", result.stderr)

    def test_group_must_reference_a_real_image(self):
        result = self.run_guard(executable(broken=True))
        self.assertNotEqual(result.returncode, 0)
        self.assertIn("missing image", result.stderr)

    def test_truncated_executable_is_rejected(self):
        self.assertNotEqual(self.run_guard(b"MZ").returncode, 0)


if __name__ == "__main__":
    unittest.main()