#!/usr/bin/env python3
"""Generate two small, distinct, valid PNG files for the POC."""
import struct
import zlib
import os

ROOT = os.path.dirname(os.path.abspath(__file__))


def chunk(tag: bytes, data: bytes) -> bytes:
    return (struct.pack(">I", len(data)) + tag + data +
            struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF))


def make_png(w: int, h: int, r: int, g: int, b: int) -> bytes:
    sig = b"\x89PNG\r\n\x1a\n"
    ihdr = chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
    raw = b""
    for _ in range(h):
        raw += b"\x00" + bytes([r, g, b] * w)
    idat = chunk(b"IDAT", zlib.compress(raw, 9))
    iend = chunk(b"IEND", b"")
    return sig + ihdr + idat + iend


a = make_png(4, 4, 255, 0, 0)        # tiny red — what the user pastes
b = make_png(32, 32, 0, 0, 255)      # larger blue — the "secret" placed after

open(os.path.join(ROOT, "png_a.png"), "wb").write(a)
open(os.path.join(ROOT, "png_b.png"), "wb").write(b)
print(f"png_a.png: {len(a)} bytes")
print(f"png_b.png: {len(b)} bytes")