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