webmail-client/backend/app/indexer/thumbs.py

"""Attachment thumbnails for the reading pane: images and a PDF's first page.

A thumbnail fits within THUMB_MAX × THUMB_MAX with its aspect ratio kept (the
browser decides how to frame it), as WebP. Made once and kept beside the blob
store — ``{BLOB_ROOT}/{user}/thumbs/{sha}-{THUMB_MAX}.webp`` — with a ``.none``
marker for files that cannot have one (a broken image, an encrypted PDF, HEIC),
so neither is attempted twice.

Attachments are hostile input: decoding runs in a worker thread, oversized
files and decompression bombs are refused.
"""

from __future__ import annotations

import io
import logging
import uuid
from collections.abc import Callable
from pathlib import Path

from app.indexer import blobstore

logger = logging.getLogger(__name__)

THUMB_MAX = 480  # px: sharp on a 2× screen for a tile of up to 230 px
_MAX_INPUT_BYTES = 40 * 1024 * 1024
_MAX_PIXELS = 64_000_000  # a larger image is refused, not decoded
_IMAGE_MIME = frozenset(
    {"image/jpeg", "image/png", "image/gif", "image/webp", "image/bmp", "image/tiff"}
)


_EXT_MIME = {
    "jpg": "image/jpeg",
    "jpeg": "image/jpeg",
    "png": "image/png",
    "gif": "image/gif",
    "webp": "image/webp",
    "bmp": "image/bmp",
    "tif": "image/tiff",
    "tiff": "image/tiff",
    "pdf": "application/pdf",
}


def thumb_type(mime_type: str, filename: str) -> str | None:
    """The type to render a thumbnail as, or None. Senders often label files
    application/octet-stream, so the extension stands in for a missing type."""
    mime = mime_type.lower()
    if mime in _IMAGE_MIME or mime == "application/pdf":
        return mime
    return _EXT_MIME.get(filename.rsplit(".", 1)[-1].lower()) if "." in filename else None


def can_thumbnail(mime_type: str) -> bool:
    return mime_type in _IMAGE_MIME or mime_type == "application/pdf"


def _path(user_id: uuid.UUID, sha: str, suffix: str) -> Path:
    return blobstore.blob_path(user_id, sha).parent.parent / "thumbs" / f"{sha}-{THUMB_MAX}{suffix}"


def _render(data: bytes, mime_type: str) -> bytes | None:
    from PIL import Image, ImageOps

    Image.MAX_IMAGE_PIXELS = _MAX_PIXELS  # beyond: DecompressionBombError
    if mime_type == "application/pdf":
        import pypdfium2 as pdfium  # type: ignore[import-untyped]

        doc = pdfium.PdfDocument(data)
        try:
            page = doc[0]
            width, height = page.get_size()
            scale = THUMB_MAX / max(width, height, 1)
            img = page.render(scale=scale, draw_annots=True).to_pil()
        finally:
            doc.close()
    else:
        img = Image.open(io.BytesIO(data))
        img.seek(0)  # an animation: its first frame
        img = ImageOps.exif_transpose(img)  # phone photos: upright
        img.thumbnail((THUMB_MAX, THUMB_MAX))
    if img.mode not in ("RGB", "RGBA"):
        img = img.convert("RGBA" if "A" in img.getbands() or img.mode == "P" else "RGB")
    out = io.BytesIO()
    img.save(out, "WEBP", quality=80, method=4)
    return out.getvalue()


def made(user_id: uuid.UUID, sha: str) -> bool:
    """Whether the attachment's thumbnail (or its absence) is settled already."""
    return _path(user_id, sha, ".webp").exists() or _path(user_id, sha, ".none").exists()


def thumbnail(
    user_id: uuid.UUID,
    sha: str,
    mime_type: str,
    size: int,
    load: Callable[[], bytes],
) -> bytes | None:
    """The thumbnail of an attachment (content hash ``sha``, ``size`` bytes,
    read with ``load`` only when it is not made yet), or None if it cannot
    have one.

    Blocking (decoding is CPU work): call it in a thread.
    """
    done = _path(user_id, sha, ".webp")
    if done.exists():
        return done.read_bytes()
    none = _path(user_id, sha, ".none")
    if none.exists() or not can_thumbnail(mime_type):
        return None
    if size > _MAX_INPUT_BYTES:
        return None
    try:
        data = load()
    except FileNotFoundError:
        return None
    try:
        thumb = _render(data, mime_type)
    except Exception as exc:  # noqa: BLE001 - any undecodable file: no thumbnail
        logger.info("no thumbnail for blob %s (%s): %s", sha, mime_type, exc)
        thumb = None
    done.parent.mkdir(parents=True, exist_ok=True)
    if thumb is None:
        none.touch()
        return None
    tmp = done.with_suffix(".tmp")
    tmp.write_bytes(thumb)
    tmp.replace(done)
    return thumb