agentic.email/backend/app/services/mime_body_extractor.py

from dataclasses import dataclass, field
from email import policy
from email.message import EmailMessage, Message
from email.parser import BytesParser
from html.parser import HTMLParser
import re

from app.services.reply_extractor import extract_new_reply_content


@dataclass
class ExtractedMimePart:
    filename: str
    content_type: str
    payload: bytes
    content_id: str | None = None
    disposition: str = "attachment"


@dataclass
class MirroredMimeContent:
    body_text: str = ""
    body_html: str | None = None
    inline_parts: list[ExtractedMimePart] = field(default_factory=list)
    file_attachments: list[ExtractedMimePart] = field(default_factory=list)


def _part_payload(part: Message) -> bytes:
    payload = part.get_payload(decode=True)
    if isinstance(payload, bytes):
        return payload
    if isinstance(payload, str):
        return payload.encode("utf-8", errors="replace")
    return b""


def _part_text(part: Message) -> str:
    try:
        value = part.get_content()
    except (LookupError, UnicodeError):
        value = _part_payload(part).decode("utf-8", errors="replace")
    return value if isinstance(value, str) else ""


def _content_id(part: Message) -> str | None:
    header = part.get("Content-ID")
    if not header:
        return None
    return header.strip().strip("<>")


def _is_attachment_part(part: Message) -> bool:
    disposition = (part.get_content_disposition() or "").lower()
    if disposition == "attachment":
        return True
    if disposition == "inline":
        return False
    filename = part.get_filename()
    content_type = part.get_content_type()
    if filename and content_type not in {"text/plain", "text/html"}:
        return True
    return False


def _walk_parts(message: Message, collector: MirroredMimeContent) -> None:
    if message.is_multipart():
        for part in message.get_payload():
            if isinstance(part, Message):
                _walk_parts(part, collector)
        return

    content_type = message.get_content_type()
    cid = _content_id(message)
    filename = message.get_filename() or (f"inline-{cid}" if cid else "attachment")

    if content_type == "text/plain" and not _is_attachment_part(message):
        text = _part_text(message).strip()
        if text and not collector.body_text:
            collector.body_text = text
        return

    if content_type == "text/html" and not _is_attachment_part(message):
        html = _part_text(message).strip()
        if html:
            collector.body_html = html
        return

    payload = _part_payload(message)
    if not payload:
        return

    extracted = ExtractedMimePart(
        filename=filename,
        content_type=content_type,
        payload=payload,
        content_id=cid,
        disposition="inline" if cid or message.get_content_disposition() == "inline" else "attachment",
    )

    if cid or extracted.disposition == "inline":
        collector.inline_parts.append(extracted)
    elif _is_attachment_part(message):
        collector.file_attachments.append(extracted)


def extract_mirror_content(raw: bytes) -> MirroredMimeContent:
    message = BytesParser(policy=policy.default).parsebytes(raw)
    collector = MirroredMimeContent()
    _walk_parts(message, collector)
    return collector


class _NewReplyHTMLParser(HTMLParser):
    QUOTE_CLASSES = {
        "gmail_attr",
        "gmail_quote",
        "protonmail_quote",
        "yahoo_quoted",
        "moz-cite-prefix",
    }
    QUOTE_IDS = {"divRplyFwdMsg", "appendonsend"}
    VOID_TAGS = {"br", "hr", "img", "meta", "link", "input"}

    def __init__(self) -> None:
        super().__init__(convert_charrefs=False)
        self.output: list[str] = []
        self.open_tags: list[str] = []
        self.stopped = False

    def _is_quote_start(self, tag: str, attrs: list[tuple[str, str | None]]) -> bool:
        values = {key.lower(): (value or "") for key, value in attrs}
        classes = set(values.get("class", "").split())
        return (
            tag.lower() == "blockquote"
            or bool(classes.intersection(self.QUOTE_CLASSES))
            or values.get("id") in self.QUOTE_IDS
        )

    def handle_starttag(self, tag: str, attrs) -> None:
        if self.stopped:
            return
        if self._is_quote_start(tag, attrs):
            self.stopped = True
            return
        self.output.append(self.get_starttag_text())
        if tag.lower() not in self.VOID_TAGS:
            self.open_tags.append(tag)

    def handle_startendtag(self, tag: str, attrs) -> None:
        if not self.stopped:
            self.output.append(self.get_starttag_text())

    def handle_endtag(self, tag: str) -> None:
        if self.stopped:
            return
        self.output.append(f"</{tag}>")
        if tag in self.open_tags:
            reverse_index = self.open_tags[::-1].index(tag)
            del self.open_tags[len(self.open_tags) - reverse_index - 1]

    def handle_data(self, data: str) -> None:
        if self.stopped:
            return
        if re.search(r"(?:^|\n)\s*On .+? wrote:\s*$", data, re.IGNORECASE):
            self.stopped = True
            return
        self.output.append(data)

    def handle_entityref(self, name: str) -> None:
        if not self.stopped:
            self.output.append(f"&{name};")

    def handle_charref(self, name: str) -> None:
        if not self.stopped:
            self.output.append(f"&#{name};")

    def html(self) -> str:
        value = "".join(self.output).strip()
        if not value:
            return ""
        return value + "".join(f"</{tag}>" for tag in reversed(self.open_tags))


def _new_reply_html(html: str | None) -> str | None:
    if not html:
        return None
    parser = _NewReplyHTMLParser()
    parser.feed(html)
    parser.close()
    value = parser.html().strip()
    return value or None


def extract_new_reply_mime_content(raw: bytes) -> MirroredMimeContent:
    """Extract only newly authored MIME content while retaining rich formatting."""
    full = extract_mirror_content(raw)
    extracted = extract_new_reply_content(
        body_text=full.body_text,
        body_html=full.body_html,
    )
    new_html = _new_reply_html(full.body_html)
    referenced_cids = {
        match.casefold()
        for match in re.findall(r"cid:([^\"'\s>)]+)", new_html or "", re.IGNORECASE)
    }
    inline_parts = [
        part
        for part in full.inline_parts
        if part.content_id and part.content_id.casefold() in referenced_cids
    ]
    return MirroredMimeContent(
        body_text=extracted["new_reply_content"].strip(),
        body_html=new_html,
        inline_parts=inline_parts,
        file_attachments=full.file_attachments,
    )