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
    # The part's position in ``msg.walk()``: where to find it again in the raw
    # message, which is the only copy kept (attachments are not stored apart).
    part_index: int = -1


@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,
        # The document <title> is not body text; unwrapped it would lead the snippet.
        clean_content_tags={"script", "style", "title"},
    )


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 index, part in enumerate(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,
                part_index=index,
            )
        )
    # An invitation sent only as the text/calendar alternative (Outlook, often) has
    # no file to show: keep it as invite.ics, so the event can be seen, answered
    # and downloaded like a Gmail one. Not when the message attaches its own .ics.
    if not any(a.filename.lower().endswith(".ics") for a in out):
        for index, part in enumerate(msg.walk()):
            if part.get_content_type() != "text/calendar" or part.get_filename():
                continue
            payload = part.get_payload(decode=True)
            if isinstance(payload, bytes) and payload.strip():
                out.append(
                    ParsedAttachment(
                        filename="invite.ics",
                        mime_type="text/calendar",
                        content_id=None,
                        payload=payload,
                        part_index=index,
                    )
                )
                break
    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),
    )


def load_message(
    raw: bytes | None = None, path: str | os.PathLike[str] | None = None
) -> EmailMessage:
    """The parsed tree of a raw message, from memory or straight from its file."""
    if path is not None:
        with open(path, "rb") as handle:
            msg = email.message_from_binary_file(handle, policy=default_policy)
    else:
        msg = email.message_from_bytes(raw or b"", policy=default_policy)
    assert isinstance(msg, EmailMessage)
    return msg


def part_payload(msg: EmailMessage, part_index: int | None, sha256: str) -> bytes | None:
    """The decoded bytes of the attachment whose content hashes to ``sha256``:
    the part at ``part_index`` when that is it, else whichever part is (rows
    stored before the index was kept). None when the message has no such part."""
    import hashlib

    parts = list(msg.walk())
    if part_index is not None and 0 <= part_index < len(parts):
        payload = parts[part_index].get_payload(decode=True)
        if isinstance(payload, bytes) and hashlib.sha256(payload).hexdigest() == sha256:
            return payload
    for part in parts:
        if part.is_multipart():
            continue
        payload = part.get_payload(decode=True)
        if isinstance(payload, bytes) and hashlib.sha256(payload).hexdigest() == sha256:
            return payload
    return None