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

"""MIME parsing via the stdlib ``email`` package with ``policy.default`` (SPEC §6.4).

Primary parsing is stdlib per SPEC §3; ``mail-parser`` is available for helpers only.
Header RFC2047 decoding is automatic under ``policy.default``.
"""

from __future__ import annotations

import email
import os
import re
from dataclasses import dataclass, field
from datetime import datetime
from email.headerregistry import AddressHeader
from email.message import EmailMessage
from email.policy import default as default_policy

import nh3

from app.threading_.normalize import normalize_message_id, normalize_reference_list

_SNIPPET_LEN = 160

# Undecodable header/body bytes come out of the stdlib parser as lone UTF-16
# surrogates (U+D800–U+DFFF). Those are invalid UTF-8, so Postgres rejects them in
# text/JSONB columns — scrub them to U+FFFD before anything is stored.
_SURROGATE_RE = re.compile("[\ud800-\udfff]")


def _scrub(value: str) -> str:
    return _SURROGATE_RE.sub("�", value) if value else value


def _scrub_addrs(addrs: list[dict[str, str]]) -> list[dict[str, str]]:
    return [{k: _scrub(v) for k, v in addr.items()} for addr in addrs]


@dataclass(frozen=True)
class ParsedAttachment:
    filename: str
    mime_type: str
    content_id: str | None
    payload: bytes


@dataclass
class ParsedEmail:
    message_id: str | None
    in_reply_to: str | None
    references: list[str]  # References + In-Reply-To, normalized, in order (§6.4)
    subject: str
    from_addr: list[dict[str, str]]
    to_addr: list[dict[str, str]]
    cc_addr: list[dict[str, str]]
    bcc_addr: list[dict[str, str]]
    date: datetime | None
    text_body: str
    html_body: str | None
    attachments: list[ParsedAttachment] = field(default_factory=list)

    @property
    def has_attachments(self) -> bool:
        # Inline resources (cid: images) carry a content_id and render in-body, not
        # as attachment chips — only real attachments flip the paperclip.
        return any(a.content_id is None for a in self.attachments)

    @property
    def snippet(self) -> str:
        return " ".join(self.text_body.split())[:_SNIPPET_LEN]


def _addresses(msg: EmailMessage, header: str) -> list[dict[str, str]]:
    value = msg[header]
    if value is None:
        return []
    out: list[dict[str, str]] = []
    if isinstance(value, AddressHeader):
        for addr in value.addresses:
            if addr.addr_spec:
                out.append({"name": str(addr.display_name or ""), "email": addr.addr_spec})
    return out


def _html_to_text(html: str) -> str:
    """Derive plain text from HTML by stripping all tags (SPEC §6.4 nh3 strip)."""
    return nh3.clean(html, tags=set(), attributes={}, strip_comments=True)


def _get_body_content(msg: EmailMessage, subtype: str) -> str | None:
    part = msg.get_body(preferencelist=(subtype,))
    if part is None:
        return None
    try:
        content = part.get_content()
    except (LookupError, ValueError, KeyError):
        payload = part.get_payload(decode=True)
        if not isinstance(payload, bytes):
            return None
        content = payload.decode("utf-8", "replace")
    return content if isinstance(content, str) else None


def _extract_attachments(msg: EmailMessage) -> list[ParsedAttachment]:
    # Walk the whole tree, not iter_attachments(): the latter is shallow and treats
    # multipart/related as a body, so an inline image nested in mixed>related>alt is
    # never yielded (its cid: reference then renders broken). A part is an attachment
    # or inline resource if it has a filename, a Content-ID, or attachment disposition;
    # the body text/plain and text/html parts have none of these and are skipped.
    out: list[ParsedAttachment] = []
    for part in msg.walk():
        if part.is_multipart() or not isinstance(part, EmailMessage):
            continue
        cid = part.get("Content-ID")
        filename = part.get_filename()
        disposition = part.get_content_disposition()
        if not (filename or cid or disposition == "attachment"):
            continue
        payload = part.get_payload(decode=True)
        if not isinstance(payload, bytes):
            continue
        # A part is an INLINE resource (cid: image rendered in-body, hidden from the
        # attachment list) only when it carries a Content-ID AND is not explicitly an
        # attachment. A part marked `Content-Disposition: attachment` is a real
        # attachment even when it also has a Content-ID (e.g. Fastmail's invite.ics),
        # so we drop the cid for it and it shows as a downloadable chip.
        inline = cid is not None and disposition != "attachment"
        out.append(
            ParsedAttachment(
                filename=_scrub(filename or "attachment"),
                mime_type=part.get_content_type(),
                content_id=cid.strip("<>") if (inline and cid is not None) else None,
                payload=payload,
            )
        )
    return out


def parse_message(raw: bytes) -> ParsedEmail:
    """Parse a raw RFC822 message (in memory) into a :class:`ParsedEmail`."""
    msg = email.message_from_bytes(raw, policy=default_policy)
    assert isinstance(msg, EmailMessage)
    return _parse_email(msg)


def parse_message_file(path: str | os.PathLike[str]) -> ParsedEmail:
    """Parse a message straight from a blob file handle (large-message path, §7)."""
    with open(path, "rb") as handle:
        msg = email.message_from_binary_file(handle, policy=default_policy)
    assert isinstance(msg, EmailMessage)
    return _parse_email(msg)


def _parse_email(msg: EmailMessage) -> ParsedEmail:
    message_id = normalize_message_id(msg.get("Message-ID"))
    in_reply_to_ids = normalize_reference_list(msg.get("In-Reply-To"))
    in_reply_to = in_reply_to_ids[0] if in_reply_to_ids else None

    # references_hdr = References then In-Reply-To message-ids, in order, deduped (§6.4).
    references: list[str] = []
    for ref in normalize_reference_list(msg.get("References")) + in_reply_to_ids:
        if ref not in references:
            references.append(ref)

    subject = str(msg.get("Subject", "") or "")

    date: datetime | None = None
    date_header = msg["Date"]
    if date_header is not None:
        try:
            date = date_header.datetime
        except (AttributeError, ValueError):
            date = None

    html_body = _get_body_content(msg, "html")
    text_body = _get_body_content(msg, "plain")
    if text_body is None and html_body is not None:
        text_body = _html_to_text(html_body)

    return ParsedEmail(
        message_id=message_id,
        in_reply_to=in_reply_to,
        references=references,
        subject=_scrub(subject),
        from_addr=_scrub_addrs(_addresses(msg, "From")),
        to_addr=_scrub_addrs(_addresses(msg, "To")),
        cc_addr=_scrub_addrs(_addresses(msg, "Cc")),
        bcc_addr=_scrub_addrs(_addresses(msg, "Bcc")),
        date=date,
        text_body=_scrub(text_body or ""),
        html_body=_scrub(html_body) if html_body is not None else None,
        attachments=_extract_attachments(msg),
    )