chorus/bot/chorus_bot/cogs/email_ingest.py

94 lines
3.4 KiB
Python

import json
import os
from pathlib import Path
import discord
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"))
# 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)
await self._poll_one(modality, host, user, password)
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")
_, data = conn.search(None, "UNSEEN")
for uid in data[0].split():
uid_str = uid.decode()
if uid_str in seen:
continue
_, msg_data = conn.fetch(uid, "(RFC822)")
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] + "..."
embed = build_ping_embed(
modality=modality, sender_id=sender, snippet=snippet,
raw_content=body, 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 _extract_body(msg) -> str:
if msg.is_multipart():
for part in msg.walk():
if part.get_content_type() == "text/plain":
return part.get_payload(decode=True).decode(errors="replace")
return "(no plain-text body found)"
return msg.get_payload(decode=True).decode(errors="replace")
async def setup(bot: commands.Bot, config):
await bot.add_cog(EmailIngestCog(bot, config))