#!/usr/bin/env python3
import os
import sys
import zipfile

LAYER_COUNT = 5
OUTER_NAME = "matryoshka.png"
FLAG_TEXT = "CTF{h3ll0_d3xt3r_m0rg4n_m4try0shk4!}"


def append_file_to_png(png_path: str, extra_path: str, out_path: str):
    with open(png_path, "rb") as f_png, open(extra_path, "rb") as f_extra, open(out_path, "wb") as f_out:
        f_out.write(f_png.read())
        f_out.write(f_extra.read())


def main():
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <folder_with_5_pngs>")
        sys.exit(1)

    input_dir = sys.argv[1]

    if not os.path.isdir(input_dir):
        print(f"[-] Not a directory: {input_dir}")
        sys.exit(1)

    # Collect PNGs and sort them for deterministic order
    pngs = [f for f in os.listdir(input_dir) if f.lower().endswith(".png")]
    pngs.sort()

    if len(pngs) < LAYER_COUNT:
        print(f"[-] Need at least {LAYER_COUNT} PNG files in {input_dir}, found {len(pngs)}")
        sys.exit(1)

    # Use exactly the first LAYER_COUNT PNGs
    pngs = pngs[:LAYER_COUNT]

    print("[+] Using these PNGs as layers (outermost -> innermost):")
    for i, name in enumerate(pngs, start=1):
        print(f"  Layer {i}: {name}")

    os.makedirs("build", exist_ok=True)

    # 1) Create innermost: flag.txt -> zipN -> imageN_with_zip
    flag_path = os.path.join("build", "flag.txt")
    with open(flag_path, "w") as f:
        f.write(FLAG_TEXT + "\n")

    # Innermost index
    innermost_idx = LAYER_COUNT - 1
    innermost_png_original = os.path.join(input_dir, pngs[innermost_idx])

    zip_path = os.path.join("build", f"layer{LAYER_COUNT}.zip")
    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as z:
        z.write(flag_path, arcname="flag.txt")

    embedded_png_path = os.path.join("build", f"layer{LAYER_COUNT}_embedded.png")
    append_file_to_png(innermost_png_original, zip_path, embedded_png_path)

    inner_file = embedded_png_path  # this will be zipped into the next outer layer

    # 2) Build outer layers (N-1 down to 1)
    for layer in range(LAYER_COUNT - 1, 0, -1):
        png_name = pngs[layer - 1]
        png_original = os.path.join(input_dir, png_name)

        # zip containing the previous (inner) embedded image
        zip_path = os.path.join("build", f"layer{layer}.zip")
        with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as z:
            # Store it under a generic name inside the zip
            z.write(inner_file, arcname=f"inner_layer.png")

        # append that zip to this layer's original PNG
        embedded_png_path = os.path.join("build", f"layer{layer}_embedded.png")
        append_file_to_png(png_original, zip_path, embedded_png_path)

        inner_file = embedded_png_path

    # 3) Copy the outermost embedded PNG as final challenge file
    with open(inner_file, "rb") as f_in, open(OUTER_NAME, "wb") as f_out:
        f_out.write(f_in.read())

    print(f"[+] Created final challenge file: {OUTER_NAME}")
    print("[+] Intermediate files are in the build/ directory.")


if __name__ == "__main__":
    main()