webmail-client/backend/app/api/messages.py

"""Message body/attachment endpoints (SPEC §9.1 Mail)."""

from __future__ import annotations

import asyncio
import contextlib
import html as html_lib
import io
import logging
import uuid
import zipfile
from typing import Annotated
from urllib.parse import quote

from fastapi import APIRouter, Depends, Query
from fastapi.responses import FileResponse, HTMLResponse, Response
from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.asyncio import AsyncSession

from app import translation
from app.api.errors import raise_api_error
from app.api.schemas import AttachmentOut
from app.api.settings import (
    add_allowed_image_sender,
    allowed_image_senders,
    blocks_trackers,
)
from app.completion import html_to_text
from app.db import get_session
from app.identity.deps import CurrentUser, require_csrf
from app.indexer import attachments as att_store
from app.indexer import rawstore, thumbs
from app.indexer.mime import parse_message
from app.indexer.sanitize import SANITIZER_VERSION, sanitize_html
from app.indexer.thumbs import thumb_type
from app.models import (
    Account,
    Attachment,
    BodyPart,
    ComposeAttachment,
    Message,
    MessageTranslation,
    User,
)
from app.translation_prep import find_names_in_signature

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/messages", tags=["mail"])

DbSession = Annotated[AsyncSession, Depends(get_session)]

# Content types safe to serve inline (SPEC §10.8); everything else is forced to
# an attachment download with a neutral type.
_INLINE_SAFE = {"image/png", "image/jpeg", "image/gif", "image/webp", "application/pdf"}
# Videos the browser can play, by extension — senders often label them
# application/octet-stream. A QuickTime .mov is usually H.264 in an MP4-style
# container: offered as video/mp4, browsers play it (an HEVC one stays unplayable,
# and the preview offers a download instead). Media cannot run script: safe inline.
_VIDEO_EXT_MIME = {
    "mp4": "video/mp4",
    "m4v": "video/mp4",
    "mov": "video/mp4",
    "webm": "video/webm",
    "ogv": "video/ogg",
}


def _inline_type(mime_type: str, filename: str) -> str | None:
    """The type to serve an attachment as for in-app viewing, or None."""
    if mime_type in _INLINE_SAFE:
        return mime_type
    ext = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
    return _VIDEO_EXT_MIME.get(ext)


# Base styles injected around the sanitized fragment so email content renders on a
# readable light canvas with images that fit the pane (the iframe is sandboxed).
_DOC_TEMPLATE = """<!doctype html><html><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<style>
  /* Self-hosted Lato so plain-text bodies use the app font (the sandboxed iframe
     can't inherit the parent's web fonts). Served from the frontend's /fonts. */
  @font-face{{font-family:"Lato";src:url("/fonts/Lato-Regular.ttf") format("truetype");
    font-weight:400;font-style:normal;font-display:swap}}
  @font-face{{font-family:"Lato";src:url("/fonts/Lato-Bold.ttf") format("truetype");
    font-weight:700;font-style:normal;font-display:swap}}
  @font-face{{font-family:"Lato";src:url("/fonts/Lato-Italic.ttf") format("truetype");
    font-weight:400;font-style:italic;font-display:swap}}
  html,body{{margin:0;padding:0}}
  body{{font-family:system-ui,-apple-system,"Segoe UI",Roboto,sans-serif;font-size:14px;
    line-height:1.5;color:#1a1a1a;background:#fff;padding:14px;word-break:break-word}}
  img{{max-width:100%;height:auto}}
  table{{max-width:100%;border-collapse:collapse}}
  blockquote{{margin:0 0 0 12px;padding-left:12px;border-left:3px solid #e0e0e0;color:#555}}
  a{{color:#2563eb}}
  pre{{white-space:pre-wrap;word-wrap:break-word}}
  /* Plain-text emails: proportional Lato instead of the monospace <pre> default. */
  .wm-plaintext{{white-space:pre-wrap;word-wrap:break-word;
    font-family:"Lato",system-ui,-apple-system,"Segoe UI",Roboto,sans-serif;
    font-size:15px;line-height:1.55}}
</style></head><body>{body}</body></html>"""


def _wrap(fragment: str) -> str:
    return _DOC_TEMPLATE.format(body=fragment)


# The sanitized message as a document of its own. The app shows it in a
# sandboxed frame (frontend MessageBody), but this URL can also be opened
# directly in a tab, on our own origin: there, this policy is what keeps any
# script the sanitizer might ever miss from running — no scripts, no plugins,
# no forms, a sandbox of its own origin, and nothing framing it elsewhere.
_MESSAGE_CSP = (
    "default-src 'none'; script-src 'none'; object-src 'none'; base-uri 'none'; "
    "form-action 'none'; frame-ancestors 'self'; img-src 'self' data: https: http:; "
    "style-src 'self' 'unsafe-inline'; font-src 'self' data: https:; media-src 'self'; "
    "sandbox allow-popups allow-popups-to-escape-sandbox"
)


def _html(document: str, trackers: int = 0, extra: dict[str, str] | None = None) -> HTMLResponse:
    # The page reads what it needs from headers: the body is the message.
    return HTMLResponse(
        content=document,
        headers={
            "Content-Security-Policy": _MESSAGE_CSP,
            "X-Trackers-Blocked": str(trackers),
            **(extra or {}),
        },
    )


async def _owned_message(db: AsyncSession, user_id: uuid.UUID, message_id: uuid.UUID) -> Message:
    msg = await db.scalar(
        select(Message).where(Message.id == message_id, Message.user_id == user_id)
    )
    if msg is None:
        raise_api_error(404, "not_found", "message not found")
    return msg


def _sender_email(message: Message) -> str:
    return message.from_addr[0].get("email", "").lower() if message.from_addr else ""


async def _render_fragment(
    db: AsyncSession, message: Message, allow_remote: bool, block: bool
) -> tuple[str, int]:
    """The sanitized body the user sees, and how many trackers came out of it.

    Re-sanitizes from the raw blob when remote images are allowed, the stored
    copy predates the current sanitizer (SPEC §10.4), or trackers are to stay
    in; otherwise the stored copy is served as is.
    """
    parts = list(await db.scalars(select(BodyPart).where(BodyPart.message_id == message.id)))
    html_part = next((p for p in parts if p.kind == "html_sanitized"), None)

    stale = html_part is not None and html_part.sanitizer_v < SANITIZER_VERSION
    # The stored copy has its trackers taken out; with blocking switched off, a
    # message that had some is rendered afresh with them left in.
    trackers_kept = not block and html_part is not None and html_part.trackers_blocked > 0
    if html_part is not None and not allow_remote and not stale and not trackers_kept:
        return html_part.content, html_part.trackers_blocked if block else 0

    raw = None
    if html_part is not None or allow_remote:  # a plain-text message has nothing to redo
        try:
            raw = await rawstore.ensure_bytes(message)
        except FileNotFoundError:
            raw = None
    parsed = parse_message(raw) if raw is not None else None
    if parsed is not None and parsed.html_body is not None:
        if stale and html_part is not None:
            # The stored copy is always the plain one: no remote images, trackers out.
            upgraded = sanitize_html(parsed.html_body, message.id)
            html_part.content = upgraded.html
            html_part.sanitizer_v = upgraded.sanitizer_v
            html_part.trackers_blocked = upgraded.trackers_blocked
            await db.commit()
            if block and not allow_remote:
                return upgraded.html, upgraded.trackers_blocked
        result = sanitize_html(
            parsed.html_body, message.id, allow_remote_images=allow_remote, block_trackers=block
        )
        return result.html, result.trackers_blocked

    if html_part is not None:
        return html_part.content, html_part.trackers_blocked
    text_part = next((p for p in parts if p.kind == "text"), None)
    body = html_lib.escape(text_part.content) if text_part else ""
    return f'<div class="wm-plaintext">{body}</div>', 0


async def _content_language(db: AsyncSession, message: Message) -> str:
    """The message's language, found (and kept) the first time it is shown."""
    if message.content_language is not None or not translation.enabled():
        return message.content_language or ""
    text = await db.scalar(
        select(BodyPart.content).where(BodyPart.message_id == message.id, BodyPart.kind == "text")
    )
    try:
        found = await translation.detect(text or message.snippet or "")
    except translation.TranslationError as exc:
        logger.info("language detection unavailable: %s", exc)
        return ""  # not kept: tried again once the service is back
    message.content_language = found
    await db.commit()
    return found


@router.get("/{message_id}/html", response_class=HTMLResponse)
async def message_html(
    message_id: uuid.UUID,
    user: CurrentUser,
    db: DbSession,
    images: bool = Query(False),
    translate: str | None = Query(None, pattern=r"^[a-z]{2,3}(-[A-Za-z]{2,4})?$"),
) -> HTMLResponse:
    """Return sanitized HTML wrapped in a document for the sandboxed iframe (SPEC §9.4, I-8).

    ``images=1`` is a one-time "Load images" — render remote images for this
    request only, without persisting the sender as trusted. ``translate=pl``
    returns it translated (app/translation.py). Headers tell the page what it
    got: X-Trackers-Blocked, X-Content-Language (the message's own language,
    when known) and, translated, X-Translated-From / X-Translated-To.
    """
    message = await _owned_message(db, user.id, message_id)
    allow_remote = images or _sender_email(message) in allowed_image_senders(user)
    block = blocks_trackers(user)
    fragment, trackers = await _render_fragment(db, message, allow_remote, block)

    headers: dict[str, str] = {}
    language = await _content_language(db, message)
    if language:
        headers["X-Content-Language"] = language
    if translate:
        fragment, source = await _translated(
            db, user, message, fragment, language, translate, allow_remote, block
        )
        headers["X-Translated-From"] = source
        headers["X-Translated-To"] = translate
    return _html(_wrap(fragment), trackers, headers)


async def _names(db: AsyncSession, user: User, message: Message) -> list[str]:
    """Names a translation leaves as they are if the model would change them:
    the people on the message (and the sender's single names, as signed:
    "Milan"), the user's own names (accounts, signatures)."""
    names: list[str] = []
    for people in (message.from_addr, message.to_addr, message.cc_addr):
        names += [p.get("name", "") for p in people if p.get("name")]
    for p in message.from_addr:
        names += [w for w in str(p.get("name", "")).split() if len(w) >= 3 and w[0].isupper()]
    rows = await db.execute(
        select(Account.display_name, Account.signature).where(Account.user_id == user.id)
    )
    for display_name, signature in rows:
        if display_name:
            names.append(display_name)
        if signature:
            names += find_names_in_signature(html_to_text(signature))
    return names


async def _translated(
    db: AsyncSession,
    user: User,
    message: Message,
    fragment: str,
    language: str,
    target: str,
    allow_remote: bool,
    block: bool,
) -> tuple[str, str]:
    """The fragment in ``target`` — from the cache, or translated and kept."""
    version = SANITIZER_VERSION * 1000 + translation.PIPELINE_VERSION
    cached = await db.get(MessageTranslation, (message.id, target, allow_remote))
    if cached is not None and cached.sanitizer_v == version:
        return cached.html, cached.source
    try:
        out = await translation.translate_html(
            fragment, language or "auto", target, names=await _names(db, user, message)
        )
    except translation.TranslationError as exc:
        raise_api_error(503, "translation_failed", str(exc))
    # The translation comes back as HTML: cleaned like any HTML we show.
    clean = sanitize_html(
        out, message.id, allow_remote_images=allow_remote, block_trackers=block
    ).html
    source = language or "auto"
    stmt = insert(MessageTranslation).values(
        message_id=message.id,
        target=target,
        images=allow_remote,
        source=source,
        html=clean,
        sanitizer_v=version,
    )
    await db.execute(
        stmt.on_conflict_do_update(
            index_elements=["message_id", "target", "images"],
            set_={"html": clean, "source": source, "sanitizer_v": version},
        )
    )
    await db.commit()
    return clean, source


@router.post("/{message_id}/allow-remote-images", dependencies=[Depends(require_csrf)])
async def allow_remote_images(
    message_id: uuid.UUID, user: CurrentUser, db: DbSession
) -> dict[str, str]:
    """Add the message's sender to the per-sender image allowlist (SPEC §9.4)."""
    message = await _owned_message(db, user.id, message_id)
    sender = _sender_email(message)
    if not sender:
        raise_api_error(400, "no_sender", "message has no sender address")
    add_allowed_image_sender(user, sender)
    await db.commit()
    return {"status": "ok", "sender": sender}


@router.get("/{message_id}/cid/{content_id}")
async def message_cid(
    message_id: uuid.UUID, content_id: str, user: CurrentUser, db: DbSession
) -> Response:
    """Serve an inline (cid:) attachment image referenced by the sanitized HTML."""
    msg = await _owned_message(db, user.id, message_id)
    att = await db.scalar(
        select(Attachment).where(
            Attachment.message_id == message_id, Attachment.content_id == content_id
        )
    )
    if att is None:
        raise_api_error(404, "not_found", "inline part not found")
    data = await _attachment_bytes(msg, att)
    media_type = att.mime_type if att.mime_type in _INLINE_SAFE else "application/octet-stream"
    return Response(content=data, media_type=media_type)


def _content_disposition(disposition: str, filename: str) -> str:
    """RFC 6266 Content-Disposition safe for non-ASCII names. HTTP header values
    are Latin-1, so a raw Polish/CJK filename (ł, ń, …) crashes the response with
    a 500; emit an ASCII fallback plus a UTF-8 ``filename*`` for modern clients."""
    ascii_name = filename.encode("ascii", "ignore").decode("ascii").replace('"', "").strip()
    fallback = ascii_name or "attachment"
    return f"{disposition}; filename=\"{fallback}\"; filename*=UTF-8''{quote(filename)}"


async def _raw_ready(msg: Message) -> None:
    """Have the message's raw source here (a hosted mailbox's is fetched from
    the mail server when not cached); a miss is left to the reader to report."""
    with contextlib.suppress(FileNotFoundError):
        await rawstore.ensure(msg)


async def _attachment_bytes(msg: Message, att: Attachment) -> bytes:
    """An attachment's bytes, taken out of its message (404 if they cannot be)."""
    await _raw_ready(msg)
    try:
        return await asyncio.to_thread(
            att_store.load, msg.user_id, msg.blob_sha256, att.blob_sha256, att.part_index
        )
    except FileNotFoundError:
        raise_api_error(404, "not_found", "attachment content missing")


async def _real_attachments(db: AsyncSession, message_id: uuid.UUID) -> list[Attachment]:
    """Downloadable (non-inline) attachments for a message, in a stable order."""
    rows = await db.scalars(
        select(Attachment)
        .where(Attachment.message_id == message_id, Attachment.content_id.is_(None))
        .order_by(Attachment.filename)
    )
    return list(rows)


@router.get("/{message_id}/attachments", response_model=list[AttachmentOut])
async def list_attachments(
    message_id: uuid.UUID, user: CurrentUser, db: DbSession
) -> list[AttachmentOut]:
    """List a message's downloadable attachments (inline cid parts excluded)."""
    await _owned_message(db, user.id, message_id)
    return [
        AttachmentOut(id=a.id, filename=a.filename, mime_type=a.mime_type, size_bytes=a.size_bytes)
        for a in await _real_attachments(db, message_id)
    ]


@router.post(
    "/{message_id}/forward-attachments",
    response_model=list[AttachmentOut],
    dependencies=[Depends(require_csrf)],
)
async def forward_attachments(
    message_id: uuid.UUID, user: CurrentUser, db: DbSession
) -> list[AttachmentOut]:
    """Attach a message's files to a forward of it: one compose attachment per
    real (non-inline) attachment. Each is taken out of the message into a blob
    of its own (shared, by hash, with any other use of the same file).
    Returns them as the composer lists its attachments."""
    msg = await _owned_message(db, user.id, message_id)
    await _raw_ready(msg)
    rows = [
        ComposeAttachment(
            user_id=user.id,
            filename=a.filename,
            mime_type=a.mime_type,
            size_bytes=a.size_bytes,
            blob_sha256=a.blob_sha256,
        )
        for a in await _real_attachments(db, message_id)
        if await asyncio.to_thread(
            att_store.materialize, user.id, msg.blob_sha256, a.blob_sha256, a.part_index
        )
    ]
    db.add_all(rows)
    await db.commit()
    return [
        AttachmentOut(id=r.id, filename=r.filename, mime_type=r.mime_type, size_bytes=r.size_bytes)
        for r in rows
    ]


@router.get("/{message_id}/attachments.zip")
async def download_attachments_zip(
    message_id: uuid.UUID,
    user: CurrentUser,
    db: DbSession,
    ids: str | None = Query(None, description="comma-separated attachment ids; all if omitted"),
) -> Response:
    """Bundle a message's attachments into a single ``Attachments.zip`` download."""
    msg = await _owned_message(db, user.id, message_id)
    wanted: set[uuid.UUID] | None = None
    if ids:
        try:
            wanted = {uuid.UUID(v) for v in ids.split(",") if v.strip()}
        except ValueError:
            raise_api_error(400, "bad_ids", "invalid attachment id list")
    attachments = [
        a for a in await _real_attachments(db, message_id) if wanted is None or a.id in wanted
    ]
    if not attachments:
        raise_api_error(404, "not_found", "no attachments to download")

    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zf:
        seen: dict[str, int] = {}
        for att in attachments:
            name = att.filename or "attachment"
            # De-dupe identical filenames within the archive (name, name-2, …).
            if name in seen:
                seen[name] += 1
                stem, dot, ext = name.rpartition(".")
                name = f"{stem}-{seen[name]}{dot}{ext}" if dot else f"{name}-{seen[name]}"
            else:
                seen[name] = 1
            zf.writestr(name, await _attachment_bytes(msg, att))

    return Response(
        content=buffer.getvalue(),
        media_type="application/zip",
        headers={"Content-Disposition": 'attachment; filename="Attachments.zip"'},
    )


@router.get("/{message_id}/attachments/{attachment_id}/thumb")
async def attachment_thumbnail(
    message_id: uuid.UUID, attachment_id: uuid.UUID, user: CurrentUser, db: DbSession
) -> Response:
    """A small preview of an image or a PDF's first page (WebP, at most
    480 px, aspect kept). 404 when the file cannot have one: the reading pane
    then shows its type icon instead."""
    msg = await _owned_message(db, user.id, message_id)
    att = await db.scalar(
        select(Attachment).where(
            Attachment.id == attachment_id, Attachment.message_id == message_id
        )
    )
    kind = thumb_type(att.mime_type, att.filename) if att is not None else None
    if att is None or kind is None:
        raise_api_error(404, "not_found", "no thumbnail")
    if not thumbs.made(msg.user_id, att.blob_sha256):
        await _raw_ready(msg)
    data = await asyncio.to_thread(
        thumbs.thumbnail,
        msg.user_id,
        att.blob_sha256,
        kind,
        att.size_bytes,
        lambda: att_store.load(msg.user_id, msg.blob_sha256, att.blob_sha256, att.part_index),
    )
    if data is None:
        raise_api_error(404, "not_found", "no thumbnail")
    # An attachment never changes under its id: the browser may keep this.
    return Response(
        content=data,
        media_type="image/webp",
        headers={"Cache-Control": "private, max-age=31536000, immutable"},
    )


@router.get("/{message_id}/attachments/{attachment_id}")
async def message_attachment(
    message_id: uuid.UUID,
    attachment_id: uuid.UUID,
    user: CurrentUser,
    db: DbSession,
    download: bool = Query(False),
) -> Response:
    """Download an attachment (inline for safe types unless ``download`` forces it)."""
    msg = await _owned_message(db, user.id, message_id)
    att = await db.scalar(
        select(Attachment).where(
            Attachment.id == attachment_id, Attachment.message_id == message_id
        )
    )
    if att is None:
        raise_api_error(404, "not_found", "attachment not found")
    if not att_store.has_file(msg.user_id, att.blob_sha256):
        await _raw_ready(msg)
    try:
        # Taken out of the message into a short-lived file (kept a day after
        # its last use), so it can be streamed with HTTP Range like before.
        path = await asyncio.to_thread(
            att_store.file, msg.user_id, msg.blob_sha256, att.blob_sha256, att.part_index
        )
    except FileNotFoundError:
        raise_api_error(404, "not_found", "attachment content missing")
    inline_type = None if download else _inline_type(att.mime_type, att.filename)
    disposition = "inline" if inline_type else "attachment"
    # Streamed from disk, with HTTP Range: a video can be played and seeked
    # without being sent whole first (Safari will not play one otherwise).
    return FileResponse(
        path,
        media_type=inline_type or "application/octet-stream",
        headers={"Content-Disposition": _content_disposition(disposition, att.filename)},
    )