chorus/bot/chorus_bot/cogs/email_ingest.py

import json
import os
from pathlib import Path
from discord.ext import commands, tasks
from imaplib import IMAP4_SSL
from email import message_from_bytes
from email.utils import parsedate_to_datetime
from chorus_bot.ping import build_ping_embed

SEEN_UIDS_PATH = Path(os.environ.get("CHORUS_SEEN_UIDS_PATH", "seen_uids.json"))

# Discord embed field values cap at 1024 chars. Truncate the raw_content we pass
# into the embed with headroom below that cap for build_ping_embed's own
# truncation marker (see ping.py). This is intentionally lossy for very long
# emails -- acceptable for MVP; the full body still lives in the source mailbox.
_RAW_CONTENT_MAX = 1000
_TRUNCATION_MARKER = "\n\n[...truncated for Discord embed limit...]"

# modality -> (imap host attr, imap user attr, imap password attr) on Config
MAILBOXES = {
    "bh_email": ("bh_email_imap_host", "bh_email_imap_user", "bh_email_imap_password"),
    "personal_email": (
        "personal_email_imap_host", "personal_email_imap_user", "personal_email_imap_password",
    ),
    "bh_text": ("google_voice_imap_host", "google_voice_imap_user", "google_voice_imap_password"),
}


def _load_seen() -> dict:
    if SEEN_UIDS_PATH.exists():
        return json.loads(SEEN_UIDS_PATH.read_text())
    return {}


def _save_seen(seen: dict) -> None:
    SEEN_UIDS_PATH.write_text(json.dumps(seen))


class EmailIngestCog(commands.Cog):
    def __init__(self, bot: commands.Bot, config):
        self.bot = bot
        self.config = config
        self.seen_uids = _load_seen()
        self.poll_mailboxes.start()

    def cog_unload(self):
        self.poll_mailboxes.cancel()

    @tasks.loop(minutes=2)
    async def poll_mailboxes(self):
        for modality, (host_attr, user_attr, pass_attr) in MAILBOXES.items():
            host = getattr(self.config, host_attr)
            user = getattr(self.config, user_attr)
            password = getattr(self.config, pass_attr)
            try:
                await self._poll_one(modality, host, user, password)
            except Exception as exc:
                print(f"Chorus: failed to poll {modality} mailbox ({host}): {exc} — will retry next cycle")

    async def _poll_one(self, modality: str, host: str, user: str, password: str):
        seen = self.seen_uids.setdefault(modality, [])
        conn = IMAP4_SSL(host)
        try:
            conn.login(user, password)
            conn.select("INBOX")
            # Use IMAP UIDs (stable across mailbox changes/expunges), not sequence
            # numbers, which shift and would make persisted seen_uids.json entries
            # refer to the wrong message later on.
            _, data = conn.uid("SEARCH", None, "UNSEEN")
            for uid in data[0].split():
                uid_str = uid.decode()
                if uid_str in seen:
                    continue
                # BODY.PEEK[] fetches without setting the \Seen flag server-side,
                # so we don't silently mark the user's real emails as read.
                _, msg_data = conn.uid("FETCH", uid, "(BODY.PEEK[])")
                raw_email = msg_data[0][1]
                msg = message_from_bytes(raw_email)
                sender = msg.get("From")
                captured_at = parsedate_to_datetime(msg.get("Date"))
                body = self._extract_body(msg)
                snippet = body if len(body) <= 120 else body[:117] + "..."
                raw_content = self._truncate_for_embed(body)

                embed = build_ping_embed(
                    modality=modality, sender_id=sender, snippet=snippet,
                    raw_content=raw_content, captured_at=captured_at,
                )
                inbox = self.bot.get_channel(self.config.inbox_channel_id)
                role_mention = f"<@&{self.config.triage_role_id}>"
                await inbox.send(content=role_mention, embed=embed)

                seen.append(uid_str)
            _save_seen(self.seen_uids)
        finally:
            conn.logout()

    @staticmethod
    def _truncate_for_embed(body: str) -> str:
        """Cap the raw_content passed into the embed at a safe length below
        Discord's 1024-char field limit. Lossy for very long emails but keeps
        ingestion from failing (and retrying forever) on oversized bodies."""
        if len(body) <= _RAW_CONTENT_MAX:
            return body
        return body[:_RAW_CONTENT_MAX] + _TRUNCATION_MARKER

    @staticmethod
    def _extract_body(msg) -> str:
        if msg.is_multipart():
            plain_part = None
            html_part = None
            for part in msg.walk():
                if part.get_content_type() == "text/plain" and plain_part is None:
                    payload = part.get_payload(decode=True)
                    if payload is not None:
                        plain_part = payload.decode(errors="replace")
                elif part.get_content_type() == "text/html" and html_part is None:
                    payload = part.get_payload(decode=True)
                    if payload is not None:
                        html_part = payload.decode(errors="replace")
            if plain_part is not None:
                return plain_part
            if html_part is not None:
                return html_part
            return "(no text content found)"
        payload = msg.get_payload(decode=True)
        if payload is None:
            return "(no text content found)"
        return payload.decode(errors="replace")


async def setup(bot: commands.Bot, config):
    await bot.add_cog(EmailIngestCog(bot, config))