#!/usr/bin/env python3
import argparse
from pathlib import Path

# Correct PNG magic number/signature
PNG_MAGIC = b"\x89PNG\r\n\x1a\n"

# A fun “corrupted” header for the Xmas CTF challenge (8 bytes)
CORRUPTED_MAGIC = b"XMASCTF!"  # exactly 8 bytes


def generate_challenge(input_png: Path, output_png: Path) -> None:
    """
    Read a valid PNG, replace the first 8 bytes (magic number) with a corrupted
    value, and write out the challenge PNG.
    """
    data = input_png.read_bytes()

    if not data.startswith(PNG_MAGIC):
        print(
            "[!] Warning: Input file does not start with a valid PNG signature. "
            "Are you sure it's a PNG?"
        )

    if len(data) < 8:
        raise ValueError("Input file is too small to be a valid PNG.")

    corrupted = CORRUPTED_MAGIC + data[8:]
    output_png.write_bytes(corrupted)

    print(f"[+] Challenge written to: {output_png}")
    print(f"[+] First 8 bytes changed from PNG magic to: {CORRUPTED_MAGIC!r}")


def solve_challenge(corrupted_png: Path, fixed_png: Path) -> None:
    """
    Read a corrupted PNG (with broken magic number), restore the correct
    PNG magic number, and write out the fixed PNG.
    """
    data = corrupted_png.read_bytes()

    if len(data) < 8:
        raise ValueError("Corrupted file is too small to be a PNG.")

    original_header = data[:8]
    fixed = PNG_MAGIC + data[8:]
    fixed_png.write_bytes(fixed)

    print(f"[+] Fixed PNG written to: {fixed_png}")
    print(f"[+] Replaced header {original_header!r} with {PNG_MAGIC!r}")


def main():
    parser = argparse.ArgumentParser(
        description="Xmas CTF: Magic Route - PNG magic number corrupter/solver."
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    # generate subcommand
    gen = subparsers.add_parser(
        "generate", help="Generate a corrupted PNG challenge from a valid PNG."
    )
    gen.add_argument("input", type=Path, help="Path to a valid input PNG.")
    gen.add_argument(
        "output",
        type=Path,
        nargs="?",
        help="Path to write corrupted challenge PNG (default: challenge.png)",
    )

    # solve subcommand
    solve = subparsers.add_parser(
        "solve",
        help="Fix a corrupted PNG (restore correct magic number).",
    )
    solve.add_argument("input", type=Path, help="Path to the corrupted PNG.")
    solve.add_argument(
        "output",
        type=Path,
        nargs="?",
        help="Path to write fixed PNG (default: fixed.png)",
    )

    args = parser.parse_args()

    if args.command == "generate":
        input_png = args.input
        output_png = args.output or Path("challenge.png")
        generate_challenge(input_png, output_png)

    elif args.command == "solve":
        corrupted_png = args.input
        fixed_png = args.output or Path("fixed.png")
        solve_challenge(corrupted_png, fixed_png)


if __name__ == "__main__":
    main()