backend/app/utils/file_validation.py

from __future__ import annotations

import hashlib
import io
from dataclasses import dataclass

from PIL import Image, UnidentifiedImageError


FORMAT_MIME_MAP = {
    "PNG": "image/png",
    "JPEG": "image/jpeg",
    "WEBP": "image/webp",
}


@dataclass(slots=True)
class ImageUploadInfo:
    width: int
    height: int
    mime_type: str
    sha256: str
    size_bytes: int
    format_name: str


class FileValidationError(ValueError):
    """Raised when an uploaded image is invalid."""


def validate_image_bytes(data: bytes, max_size_mb: int) -> ImageUploadInfo:
    size_bytes = len(data)
    max_size_bytes = max_size_mb * 1024 * 1024
    if size_bytes == 0:
        raise FileValidationError("Uploaded image is empty.")
    if size_bytes > max_size_bytes:
        raise FileValidationError(f"Image exceeds the {max_size_mb} MB upload limit.")

    try:
        image = Image.open(io.BytesIO(data))
        image.verify()
        image = Image.open(io.BytesIO(data))
    except (UnidentifiedImageError, OSError) as exc:
        raise FileValidationError("Unsupported or corrupted image file.") from exc

    format_name = (image.format or "").upper()
    mime_type = FORMAT_MIME_MAP.get(format_name)
    if mime_type is None:
        raise FileValidationError("Only PNG, JPEG, and WEBP images are supported.")

    sha256 = hashlib.sha256(data).hexdigest()
    return ImageUploadInfo(
        width=image.width,
        height=image.height,
        mime_type=mime_type,
        sha256=sha256,
        size_bytes=size_bytes,
        format_name=format_name.lower(),
    )