webmail-client/backend/app/outbox/inline_images.py

"""Inline images in an outgoing message's HTML.

A reply quotes the original as the reading pane shows it, so an image embedded
in the original (``cid:`` — a signature logo, a pasted screenshot) appears in
the quote as this server's own URL, ``/api/messages/{id}/cid/{content-id}``.
Sent as is, that URL means nothing to the recipient (and to anyone else,
including ourselves once the original is gone). Before sending, each such image
is put back into the message as an inline part with a new Content-ID — as any
mail program does — and one that can no longer be had is left out.

An image written into the HTML itself (``data:image/…;base64,…`` — a signature
logo pasted into the editor) goes the same way: Gmail and other mail programs
do not show ``data:`` images in mail they receive, so each becomes an inline
part too (one per distinct image).
"""

from __future__ import annotations

import asyncio
import base64
import binascii
import hashlib
import re
import uuid
from dataclasses import dataclass
from email.utils import make_msgid
from urllib.parse import unquote

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.indexer import attachments as att_store
from app.indexer import rawstore
from app.indexer.sanitize import INTERNAL_CID
from app.models import Attachment, Message

_DATA_IMG = re.compile(r"^data:(image/[a-z0-9.+-]+);base64,(.*)$", re.IGNORECASE | re.DOTALL)
_EXT = {"jpeg": "jpg", "svg+xml": "svg"}

_IMG = re.compile(r"<img\b[^>]*?\bsrc\s*=\s*([\"'])(.*?)\1[^>]*>", re.IGNORECASE | re.DOTALL)


@dataclass(frozen=True)
class InlinePart:
    cid: str  # without angle brackets
    mime_type: str
    filename: str
    data: bytes


def data_image(src: str, domain: str) -> InlinePart | None:
    """The inline part for a ``data:image/…;base64`` URL, or None if it is not one."""
    m = _DATA_IMG.match(src.strip())
    if m is None:
        return None
    try:
        data = base64.b64decode(re.sub(r"\s+", "", m.group(2)), validate=True)
    except (binascii.Error, ValueError):
        return None
    if not data:
        return None
    mime_type = m.group(1).lower()
    subtype = mime_type.split("/", 1)[1]
    name = f"image.{_EXT.get(subtype, subtype)}"
    return InlinePart(make_msgid(domain=domain)[1:-1], mime_type, name, data)


def internal_cid(src: str) -> tuple[uuid.UUID, str] | None:
    """The (message id, content id) an image URL of this server points at."""
    value = src.strip()
    if value.lower().startswith(("cid:", "data:")):
        return None
    m = INTERNAL_CID.search(value)
    if m is None:
        return None
    try:
        return uuid.UUID(m.group(1)), unquote(m.group(2))
    except ValueError:
        return None


async def embed(
    session: AsyncSession, user_id: uuid.UUID, html: str, domain: str
) -> tuple[str, list[InlinePart]]:
    """``html`` with this server's image URLs turned into ``cid:`` references,
    and the parts to send them in."""
    parts: dict[tuple[uuid.UUID, str], InlinePart | None] = {}
    # Images written into the HTML, by content: the same logo twice is one part.
    data_parts: dict[str, InlinePart | None] = {}

    def data_key(src: str) -> str:
        return hashlib.sha256(src.strip().encode()).hexdigest()

    for match in _IMG.finditer(html):
        src = match.group(2)
        if src.strip().lower().startswith("data:"):
            key = data_key(src)
            if key not in data_parts:
                data_parts[key] = data_image(src, domain)
            continue
        target = internal_cid(src)
        if target is None or target in parts:
            continue
        message_id, content_id = target
        found = (
            await session.execute(
                select(Attachment, Message)
                .join(Message, Message.id == Attachment.message_id)
                .where(
                    Attachment.message_id == message_id,
                    Attachment.content_id == content_id,
                    Message.user_id == user_id,
                )
                .limit(1)
            )
        ).first()
        part: InlinePart | None = None
        if found is not None:
            att, message = found
            try:
                await rawstore.ensure(message)  # a hosted mailbox's: fetched if need be
                data = await asyncio.to_thread(
                    att_store.load, user_id, message.blob_sha256, att.blob_sha256, att.part_index
                )
                cid = make_msgid(domain=domain)[1:-1]
                part = InlinePart(cid, att.mime_type, att.filename, data)
            except FileNotFoundError:
                part = None
        parts[target] = part

    def replace(match: re.Match[str]) -> str:
        quote = match.group(1)
        src = match.group(2)
        if src.strip().lower().startswith("data:"):
            data_part = data_parts.get(data_key(src))
            if data_part is None:
                return match.group(0)  # not an image we can carry: left as it was
            return match.group(0).replace(
                f"src={quote}{src}{quote}", f"src={quote}cid:{data_part.cid}{quote}", 1
            )
        target = internal_cid(match.group(2))
        if target is None:
            return match.group(0)
        part = parts.get(target)
        if part is None:
            return ""  # gone: no broken image for the recipient
        return match.group(0).replace(
            f"src={quote}{match.group(2)}{quote}", f"src={quote}cid:{part.cid}{quote}", 1
        )

    out = _IMG.sub(replace, html) if parts or data_parts else html
    quoted = [p for p in parts.values() if p is not None]
    return out, quoted + [p for p in data_parts.values() if p is not None]