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

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

from __future__ import annotations

import html as html_lib
import io
import uuid
import zipfile
from typing import Annotated

from fastapi import APIRouter, Depends, Query
from fastapi.responses import 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
from app.indexer.mime import parse_message
from app.indexer.sanitize import SANITIZER_VERSION, sanitize_html
from app.models import Attachment, BodyPart, 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"}

# 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>
  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}}
</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) -> 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).
    """
    message = await _owned_message(db, user.id, message_id)
    allow_remote = _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"<pre>{body}</pre>"))


@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)


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.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}")
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")
    data = blobstore.load(msg.user_id, att.blob_sha256)
    inline = att.mime_type in _INLINE_SAFE and not download
    media_type = att.mime_type if inline else "application/octet-stream"
    disposition = "inline" if inline else "attachment"
    return Response(
        content=data,
        media_type=media_type,
        headers={"Content-Disposition": f'{disposition}; filename="{att.filename}"'},
    )