timmy-talking-turd/scripts/reencode_image.py

#!/usr/bin/env python3
"""Re-encode one image safely: verify decode, strip metadata, bound dimensions.

Fixed argv contract only:
    reencode_image.py --in SOURCE --out TARGET [--max-bytes N]
                      [--max-dimension N] [--max-pixels N]

Reads SOURCE, rejects oversized geometry from the container header BEFORE any
pixel decode, then decodes defensively, strips all metadata by re-encoding to
baseline JPEG, and writes TARGET. Prints one JSON line on success. Never prints
image bytes or base64 to stdout/stderr.

Exit codes are a contract the caller maps to client-visible outcomes:
    0  success
    2  rejected hostile/invalid input (client's fault)
    3  processing unavailable (server runtime fault, not the client's)
"""
import argparse
import io
import json
import os
import sys
import warnings

EXIT_OK = 0
EXIT_REJECTED = 2
EXIT_UNAVAILABLE = 3

# Pinned, immutable production image runtime. This is the exact Python/Pillow
# toolchain the service re-encoder is built and tested against; the deployment
# smoke gate (scripts/deploy_staging.py and CI) verifies the provisioned
# runtime matches this record before promoting. Changing the runtime must bump
# this pin in lockstep with requirements-test.txt (Pillow==12.3.0) and the
# re-encode resource contract.
PINNED_PYTHON_VERSION = "3.11"
PINNED_PILLOW_VERSION = "12.3.0"

# Default total-pixel ceiling. Bounded independently of the per-dimension cap so
# wide-and-short or tall-and-narrow geometry cannot smuggle a huge pixel budget
# past a per-side check. 4096*4096 matches the accepted dimension envelope.
DEFAULT_MAX_PIXELS = 4096 * 4096


def parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser()
    p.add_argument("--in", dest="source")
    p.add_argument("--out", dest="target")
    p.add_argument("--max-bytes", type=int, default=4 * 1024 * 1024)
    p.add_argument("--max-dimension", type=int, default=4096)
    p.add_argument("--max-pixels", type=int, default=DEFAULT_MAX_PIXELS)
    p.add_argument("--verify-pin", action="store_true",
                   help="report the actual Python/Pillow runtime and exit; refuse "
                        "if it does not match the pinned immutable production runtime")
    return p


def _emit_pin(pillow_version: str) -> int:
    """Verify the provisioned runtime matches the pinned immutable toolchain."""
    actual_python = f"{sys.version_info.major}.{sys.version_info.minor}"
    matched = (actual_python == PINNED_PYTHON_VERSION and pillow_version == PINNED_PILLOW_VERSION)
    print(json.dumps({
        "ok": matched,
        "python": actual_python,
        "pillow": pillow_version,
        "pinned": {"python": PINNED_PYTHON_VERSION, "pillow": PINNED_PILLOW_VERSION},
    }))
    # A runtime that is not the pinned production toolchain is a server-side
    # provisioning fault, not client input: map to processing-unavailable.
    return EXIT_OK if matched else EXIT_UNAVAILABLE


def reject(error: str) -> int:
    print(json.dumps({"ok": False, "error": error}))
    return EXIT_REJECTED


def main() -> int:
    args = parser().parse_args()

    try:
        from PIL import Image as _Image
    except Exception:
        # Server-side runtime fault: the caller must map this to
        # processing-unavailable, never to "your photo is corrupt".
        print(json.dumps({"ok": False, "error": "unavailable"}))
        return EXIT_UNAVAILABLE

    # Pin verification mode: the deployment smoke gate invokes the re-encoder
    # with --verify-pin to prove the provisioned, immutable runtime is exactly
    # the one the service was built and tested against. No live host change.
    if args.verify_pin:
        return _emit_pin(getattr(_Image, "__version__", "unknown"))

    # Decompression-bomb warnings are errors here. Pillow's default threshold
    # only warns and then hands back a fully decoded image, which is not a
    # rejection; promote it so any bomb path raises instead.
    warnings.simplefilter("error", _Image.DecompressionBombWarning)
    # Refuse to let Pillow allocate beyond our own accepted pixel envelope.
    _Image.MAX_IMAGE_PIXELS = max(1, args.max_pixels)

    try:
        source_size = os.path.getsize(args.source)
    except OSError:
        return reject("rejected")
    if source_size <= 0 or source_size > args.max_bytes:
        return reject("rejected")

    # --- Header-only geometry gate, before any pixel decode. -----------------
    # Image.open() parses the header lazily, so probe.size is available without
    # allocating the pixel buffer. Rejecting here keeps a 12000x12000 bomb at
    # header cost instead of ~163 MiB of decoded pixels.
    try:
        with _Image.open(args.source) as probe:
            width, height = probe.size
    except _Image.DecompressionBombWarning:
        return reject("dimensions")
    except _Image.DecompressionBombError:
        return reject("dimensions")
    except Exception:
        return reject("malformed")

    if width <= 0 or height <= 0:
        return reject("malformed")
    if width > args.max_dimension or height > args.max_dimension:
        return reject("dimensions")
    if width * height > args.max_pixels:
        return reject("dimensions")

    # --- Only now is a full decode allowed. ---------------------------------
    try:
        with _Image.open(args.source) as probe:
            probe.load()
            # Re-check after decode: a hostile container can declare small
            # geometry in its header and expand during decode.
            if probe.size[0] > args.max_dimension or probe.size[1] > args.max_dimension:
                return reject("dimensions")
            if probe.size[0] * probe.size[1] > args.max_pixels:
                return reject("dimensions")
            image = probe.convert("RGB")
    except (_Image.DecompressionBombWarning, _Image.DecompressionBombError):
        return reject("dimensions")
    except Exception:
        return reject("malformed")

    buffer = io.BytesIO()
    try:
        # Baseline JPEG re-encode drops EXIF/GPS/XMP/tEXt entirely.
        image.save(buffer, "JPEG", quality=85, optimize=True, progressive=False)
    except Exception:
        return reject("encode")

    data = buffer.getvalue()
    if not data or data[:3] != b"\xff\xd8\xff" or len(data) > args.max_bytes:
        return reject("encode")

    try:
        with open(args.target, "wb") as handle:
            handle.write(data)
    except OSError:
        # Cannot write the private temp target: server-side fault.
        print(json.dumps({"ok": False, "error": "unavailable"}))
        return EXIT_UNAVAILABLE

    print(json.dumps({
        "ok": True,
        "format": "jpeg",
        "width": image.size[0],
        "height": image.size[1],
        "bytes": len(data),
        "metadataStripped": True,
    }))
    return EXIT_OK


if __name__ == "__main__":
    sys.exit(main())