app/image_guard.py

"""Untrusted image upload validation and normalisation.

Security contract (see docs/SECURITY.md):
  * hard byte cap enforced before decoding;
  * content sniffing with Pillow -- the client-supplied filename, extension
    and Content-Type are never trusted;
  * decompression-bomb guards (Pillow ``MAX_IMAGE_PIXELS`` + explicit pixel
    and dimension caps);
  * SVG rejected outright (scriptable XML, not an image for our purposes);
  * output is a freshly re-encoded RGB PNG (EXIF/ICC/XMP metadata stripped,
    orientation applied), downscaled so the pure-Python per-pixel PNG readers
    in ``forge/`` and the LLM image payload stay cheap.
"""

from __future__ import annotations

import io
from dataclasses import dataclass

from PIL import Image, ImageOps

# Pillow's own bomb guard; load_settings() re-applies the configured value
# at import time of the pipeline (module-level default keeps tests simple).
Image.MAX_IMAGE_PIXELS = 40_000_000

ALLOWED_FORMATS = {"PNG", "JPEG", "WEBP", "GIF", "BMP"}


class ImageRejected(Exception):
    """Raised when an upload fails validation. Carries a user-safe reason."""

    def __init__(self, reason: str, code: str = "image_rejected") -> None:
        super().__init__(reason)
        self.reason = reason
        self.code = code


@dataclass(frozen=True)
class NormalizedImage:
    png_bytes: bytes
    width: int
    height: int
    original_format: str
    original_width: int
    original_height: int
    downscaled: bool


def _sniff(data: bytes) -> Image.Image:
    """Open enough of the image to identify its format and dimensions.

    Do not decode pixels here.  Callers must enforce declared dimension and
    pixel caps before invoking ``load()``; that ordering is the protection
    against small compressed files expanding into excessive work or memory.
    """
    if not data:
        raise ImageRejected("The uploaded file is empty.", "empty_file")
    # Cheap pre-check: SVG (and other XML) is rejected before Pillow ever
    # sees it -- it is a document, not a raster image.
    head = data[:512].lstrip().lower()
    if head.startswith(b"<") and (b"<svg" in head or b"<?xml" in head):
        raise ImageRejected(
            "SVG and other vector/XML documents are not accepted. "
            "Please upload a raster image (PNG, JPEG, WebP, GIF or BMP).",
            "svg_rejected",
        )
    try:
        img = Image.open(io.BytesIO(data))
    except Image.DecompressionBombError as exc:
        raise ImageRejected(
            "The image is unreasonably large (decompression guard).",
            "image_too_large",
        ) from exc
    except Exception as exc:  # UnidentifiedImageError and friends
        raise ImageRejected(
            "The uploaded file could not be decoded as an image. "
            "Accepted formats: PNG, JPEG, WebP, GIF, BMP.",
            "not_an_image",
        ) from exc
    fmt = (img.format or "").upper()
    if fmt not in ALLOWED_FORMATS:
        img.close()
        raise ImageRejected(
            f"Unsupported image format {fmt or 'unknown'!r}. "
            "Accepted formats: PNG, JPEG, WebP, GIF, BMP.",
            "unsupported_format",
        )
    return img


def validate_and_normalize(
    data: bytes,
    *,
    max_bytes: int = 10 * 1024 * 1024,
    max_pixels: int = 40_000_000,
    max_side: int = 8192,
    normalize_max_side: int = 1024,
) -> NormalizedImage:
    """Validate raw upload bytes and return a normalised RGB PNG.

    Raises ImageRejected with a user-safe reason on any violation.
    """
    if len(data) > max_bytes:
        raise ImageRejected(
            f"The file is {len(data) / (1024 * 1024):.1f} MiB; the limit is "
            f"{max_bytes // (1024 * 1024)} MiB.",
            "file_too_large",
        )

    Image.MAX_IMAGE_PIXELS = max_pixels
    img = _sniff(data)
    original_format = (img.format or "").upper()
    original_width, original_height = img.size

    if original_width <= 0 or original_height <= 0:
        img.close()
        raise ImageRejected(
            "The image has invalid dimensions.",
            "image_too_large",
        )
    if original_width * original_height > max_pixels:
        img.close()
        raise ImageRejected(
            f"The image has {original_width * original_height:,} pixels; the limit is "
            f"{max_pixels:,}.",
            "image_too_large",
        )
    if max(original_width, original_height) > max_side:
        img.close()
        raise ImageRejected(
            f"The image is {original_width}x{original_height}px; the longest side may "
            f"be at most {max_side}px.",
            "image_too_large",
        )

    source = img
    try:
        # Pixel decode happens only after the cheap header-based limits above.
        try:
            img.load()
        except Image.DecompressionBombError as exc:
            raise ImageRejected(
                "The image is unreasonably large (decompression guard).",
                "image_too_large",
            ) from exc
        except Exception as exc:
            raise ImageRejected(
                "The uploaded file could not be fully decoded as an image.",
                "not_an_image",
            ) from exc

        # Apply EXIF orientation, then flatten to RGB (drops alpha compositing
        # surprises and strips all metadata when re-encoded).
        img = ImageOps.exif_transpose(img)
        if img.mode == "P":
            img = img.convert("RGBA")
        if img.mode in ("RGBA", "LA"):
            background = Image.new("RGB", img.size, (255, 255, 255))
            alpha = img.getchannel("A") if "A" in img.getbands() else None
            background.paste(img.convert("RGB"), mask=alpha)
            img = background
        elif img.mode != "RGB":
            img = img.convert("RGB")

        downscaled = False
        if max(img.size) > normalize_max_side:
            img = ImageOps.contain(img, (normalize_max_side, normalize_max_side))
            downscaled = True

        buf = io.BytesIO()
        img.save(buf, format="PNG", optimize=True)
        return NormalizedImage(
            png_bytes=buf.getvalue(),
            width=img.size[0],
            height=img.size[1],
            original_format=original_format,
            original_width=original_width,
            original_height=original_height,
            downscaled=downscaled,
        )
    finally:
        source.close()
        if img is not source:
            img.close()