106 lines
3.1 KiB
Python
106 lines
3.1 KiB
Python
from dataclasses import dataclass, field
|
|
from email import policy
|
|
from email.message import EmailMessage, Message
|
|
from email.parser import BytesParser
|
|
|
|
|
|
@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
|