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

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

from __future__ import annotations

import asyncio
import html as html_lib
import io
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.ext.asyncio import AsyncSession

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
from app.db import get_session
from app.identity.deps import CurrentUser, require_csrf
from app.indexer import blobstore, 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 Attachment, BodyPart, ComposeAttachment, Message

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)


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 ""


@router.get("/{message_id}/html", response_class=HTMLResponse)
async def message_html(
    message_id: uuid.UUID,
    user: CurrentUser,
    db: DbSession,
    images: bool = Query(False),
) -> HTMLResponse:
    """Return sanitized HTML wrapped in a document for the sandboxed iframe (SPEC §9.4, I-8).

    Re-sanitizes from the raw blob when the sender is image-approved, or when the
    stored copy predates the current sanitizer version (SPEC §10.4). ``images=1``
    is a one-time "Load images" — render remote images for this request only,
    without persisting the sender as trusted.
    """
    message = await _owned_message(db, user.id, message_id)
    allow_remote = images or _sender_email(message) in allowed_image_senders(user)

    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
    if html_part is not None and not allow_remote and not stale:
        return HTMLResponse(content=_wrap(html_part.content))

    # Re-sanitize the original HTML (image-approved sender, or a stale stored copy).
    # If the raw blob is unavailable, fall back to whatever is already stored.
    if allow_remote or stale:
        try:
            raw = blobstore.load(user.id, message.blob_sha256)
        except FileNotFoundError:
            raw = None
        if raw is not None:
            parsed = parse_message(raw)
            if parsed.html_body is not None:
                result = sanitize_html(
                    parsed.html_body, message.id, allow_remote_images=allow_remote
                )
                if html_part is not None and not allow_remote:
                    html_part.content = result.html  # persist the upgraded sanitization
                    html_part.sanitizer_v = result.sanitizer_v
                    await db.commit()
                return HTMLResponse(content=_wrap(result.html))

    if html_part is not None:
        return HTMLResponse(content=_wrap(html_part.content))
    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 HTMLResponse(content=_wrap(f'<div class="wm-plaintext">{body}</div>'))


@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 = blobstore.load(msg.user_id, att.blob_sha256)
    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 _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, sharing the stored blob — nothing is copied.
    Returns them as the composer lists its attachments."""
    await _owned_message(db, user.id, message_id)
    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 blobstore.exists(user.id, a.blob_sha256)
    ]
    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, blobstore.load(msg.user_id, att.blob_sha256))

    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")
    data = await asyncio.to_thread(thumbs.thumbnail, msg.user_id, att.blob_sha256, kind)
    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")
    path = blobstore.blob_path(msg.user_id, att.blob_sha256)
    if not path.exists():
        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)},
    )