#!/usr/bin/env python3
"""
make_test_image.py — Generate a synthetic disk image for carving exercises.

Embeds a JPEG, PNG, PDF, and ZIP into a binary blob padded with random noise.
Prints the offset of each embedded file so students can verify their carver output.

Usage:
    python3 make_test_image.py [output_file]   (default: test.img)
"""

import io
import os
import random
import struct
import sys
import zlib
import zipfile


def make_jpeg() -> bytes:
    """Return a minimal valid 1×1 white JPEG."""
    return bytes([
        0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10,
        0x4A, 0x46, 0x49, 0x46, 0x00,
        0x01, 0x01, 0x00, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00,
        0xFF, 0xDB, 0x00, 0x43, 0x00,
        *([8] * 64),
        0xFF, 0xC0, 0x00, 0x0B,
        0x08, 0x00, 0x01, 0x00, 0x01, 0x01, 0x01, 0x11, 0x00,
        0xFF, 0xC4, 0x00, 0x1F, 0x00,
        0x00, 0x01, 0x05, 0x01, 0x01, 0x01, 0x01, 0x01,
        0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
        0x08, 0x09, 0x0A, 0x0B,
        0xFF, 0xDA, 0x00, 0x08,
        0x01, 0x01, 0x00, 0x00, 0x3F, 0x00,
        0xFB, 0x28, 0xA2, 0x8A,
        0xFF, 0xD9,
    ])


def make_png() -> bytes:
    """Return a minimal valid 1×1 red PNG."""
    def chunk(name: bytes, data: bytes) -> bytes:
        crc = zlib.crc32(name + data) & 0xFFFFFFFF
        return struct.pack('>I', len(data)) + name + data + struct.pack('>I', crc)

    ihdr = struct.pack('>IIBBBBB', 1, 1, 8, 2, 0, 0, 0)
    idat = zlib.compress(b'\x00\xff\x00\x00')
    return (b'\x89PNG\r\n\x1a\n'
            + chunk(b'IHDR', ihdr)
            + chunk(b'IDAT', idat)
            + chunk(b'IEND', b''))


def make_pdf() -> bytes:
    """Return a minimal structurally valid PDF with xref table and startxref."""
    buf = bytearray()
    offsets = []

    buf += b'%PDF-1.4\n'

    offsets.append(len(buf))
    buf += b'1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n'

    offsets.append(len(buf))
    buf += b'2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n'

    offsets.append(len(buf))
    buf += b'3 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>\nendobj\n'

    xref_offset = len(buf)
    buf += b'xref\n0 4\n'
    buf += b'0000000000 65535 f \n'
    for off in offsets:
        buf += f'{off:010d} 00000 n \n'.encode()

    buf += b'trailer\n<< /Size 4 /Root 1 0 R >>\n'
    buf += b'startxref\n'
    buf += f'{xref_offset}\n'.encode()
    buf += b'%%EOF\n'

    return bytes(buf)


def make_zip() -> bytes:
    """Return a ZIP archive containing a short text file."""
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, 'w', zipfile.ZIP_STORED) as zf:
        zf.writestr('hello.txt', 'hello world')
    return buf.getvalue()


def build_image(seed: int | None = None) -> tuple[bytes, dict[str, int]]:
    """
    Assemble the test image.

    File order, leading gap, inter-file gaps, and trailing gap are all chosen
    randomly so each run produces a different layout.  Pass `seed` for a
    reproducible image (useful for answer keys).

    Returns (image_bytes, {format: offset}) so callers can print the manifest.
    """
    rng = random.Random(seed)

    makers = [('jpg', make_jpeg), ('png', make_png),
              ('pdf', make_pdf),  ('zip', make_zip)]
    rng.shuffle(makers)                         # randomise file order

    offsets = {}
    buf     = bytearray(os.urandom(rng.randint(64, 1024)))  # leading noise

    for i, (name, make) in enumerate(makers):
        offsets[name] = len(buf)
        buf += make()
        if i < len(makers) - 1:
            buf += os.urandom(rng.randint(64, 2048))        # inter-file noise

    buf += os.urandom(rng.randint(64, 512))                 # trailing noise
    return bytes(buf), offsets


def main() -> None:
    import argparse
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('output', nargs='?', default='test.img',
                    help='Output file (default: test.img)')
    ap.add_argument('--seed', type=int, default=None,
                    help='RNG seed for a reproducible layout (default: random)')
    args = ap.parse_args()

    out  = args.output
    seed = args.seed

    image, offsets = build_image(seed=seed)

    with open(out, 'wb') as f:
        f.write(image)

    print(f"wrote {out}  ({len(image):,} bytes)")
    print()
    print(f"  {'Format':<8} {'Offset (hex)':<16} {'Offset (dec)':<14} Description")
    print(f"  {'─'*60}")
    descs = {'jpg': 'JPEG image (FF D8 FF)',
             'png': 'PNG image (89 50 4E 47)',
             'pdf': 'PDF document (25 50 44 46)',
             'zip': 'ZIP archive (50 4B 03 04)'}
    for fmt, off in offsets.items():
        print(f"  {fmt:<8} {off:#016x}   {off:<14} {descs[fmt]}")
    if seed is not None:
        print(f"  (seed {seed} — re-run with --seed {seed} to reproduce this layout)")
    print()
    print("Run:  python3 carve.py -i", out, "-o carved/")


if __name__ == '__main__':
    main()