myagents/myclaude/mail/client.py

"""IMAP and SMTP client wrappers."""

import imaplib
import smtplib
from email.message import EmailMessage
from email.utils import formataddr
from typing import Any

from myclaude.mail.parser import extract_attachments, extract_text
from myclaude.mail.utils import decode_header_value, format_date


class IMAPClient:
    """Simple IMAP client wrapper."""

    def __init__(
        self,
        host: str,
        port: int,
        email: str,
        password: str,
    ) -> None:
        self.host = host
        self.port = port
        self.email = email
        self.password = password
        self._conn: imaplib.IMAP4_SSL | None = None

    def connect(self) -> "IMAPClient":
        self._conn = imaplib.IMAP4_SSL(self.host, self.port)
        self._conn.login(self.email, self.password)
        return self

    def disconnect(self) -> None:
        if self._conn:
            try:
                self._conn.close()
                self._conn.logout()
            except Exception:
                pass
            self._conn = None

    def __enter__(self) -> "IMAPClient":
        return self.connect()

    def __exit__(self, *_: Any) -> None:
        self.disconnect()

    def inbox_summaries(self, limit: int = 20) -> list[dict[str, Any]]:
        """Fetch recent inbox message summaries."""
        assert self._conn is not None
        status, _ = self._conn.select("INBOX")
        if status != "OK":
            raise RuntimeError("Failed to select INBOX")

        _, data = self._conn.search(None, "ALL")
        msg_ids = data[0].split()
        recent_ids = msg_ids[-limit:]
        recent_ids.reverse()

        summaries: list[dict[str, Any]] = []
        for msg_id in recent_ids:
            msg_id_str = msg_id.decode("ascii")
            _, fetched = self._conn.fetch(
                msg_id_str,
                "(BODY.PEEK[HEADER.FIELDS (FROM SUBJECT DATE)])",
            )
            raw_headers = b""
            for item in fetched:
                if isinstance(item, tuple):
                    raw_headers += item[1]

            from email import message_from_bytes

            header_msg = message_from_bytes(raw_headers)
            summaries.append(
                {
                    "uid": msg_id_str,
                    "from": decode_header_value(header_msg.get("From")),
                    "subject": decode_header_value(header_msg.get("Subject")),
                    "date": format_date(header_msg.get("Date")),
                }
            )
        return summaries

    def fetch_raw(self, uid: str) -> bytes:
        """Fetch raw RFC822 bytes by UID."""
        assert self._conn is not None
        status, _ = self._conn.select("INBOX")
        if status != "OK":
            raise RuntimeError("Failed to select INBOX")
        _, data = self._conn.fetch(uid, "(RFC822)")
        raw = b""
        for item in data:
            if isinstance(item, tuple):
                raw += item[1]
        return raw

    def fetch_message(self, uid: str) -> dict[str, Any]:
        """Fetch full message by UID."""
        raw = self.fetch_raw(uid)

        from email import message_from_bytes

        msg = message_from_bytes(raw)
        plain, html = extract_text(msg)
        attachments = extract_attachments(msg)

        return {
            "uid": uid,
            "from": decode_header_value(msg.get("From")),
            "to": decode_header_value(msg.get("To")),
            "cc": decode_header_value(msg.get("Cc")),
            "subject": decode_header_value(msg.get("Subject")),
            "date": format_date(msg.get("Date")),
            "message_id": decode_header_value(msg.get("Message-ID")),
            "reply_to": decode_header_value(msg.get("Reply-To")),
            "references": decode_header_value(msg.get("References")),
            "plain": plain,
            "html": html,
            "attachments": attachments,
        }


class SMTPClient:
    """Simple SMTP client wrapper."""

    def __init__(
        self,
        host: str,
        port: int,
        email: str,
        password: str,
    ) -> None:
        self.host = host
        self.port = port
        self.email = email
        self.password = password

    def send(
        self,
        to_addrs: list[str],
        subject: str,
        body: str,
        cc_addrs: list[str] | None = None,
        bcc_addrs: list[str] | None = None,
        attachments: list[str] | None = None,
        is_html: bool = False,
    ) -> None:
        """Send an email message."""
        msg = EmailMessage()
        msg["From"] = formataddr((self.email.split("@")[0], self.email))
        msg["To"] = ", ".join(to_addrs)
        if cc_addrs:
            msg["Cc"] = ", ".join(cc_addrs)
        msg["Subject"] = subject

        content_type = "text/html" if is_html else "text/plain"
        msg.set_content(body, subtype=content_type.split("/")[1])

        if attachments:
            for path in attachments:
                from pathlib import Path

                p = Path(path)
                if not p.is_file():
                    raise FileNotFoundError(f"Attachment not found: {path}")
                data = p.read_bytes()
                msg.add_attachment(
                    data,
                    maintype="application",
                    subtype="octet-stream",
                    filename=p.name,
                )

        all_recipients = list(to_addrs)
        if cc_addrs:
            all_recipients.extend(cc_addrs)
        if bcc_addrs:
            all_recipients.extend(bcc_addrs)

        with smtplib.SMTP_SSL(self.host, self.port) as server:
            server.login(self.email, self.password)
            server.sendmail(self.email, all_recipients, msg.as_string())

    def send_reply(
        self,
        original: dict[str, Any],
        body: str,
        is_html: bool = False,
    ) -> None:
        """Send a reply to an existing message."""
        to_addr = original.get("reply_to") or original["from"]
        if not to_addr:
            raise ValueError("Original message has no From or Reply-To address")

        subject = original.get("subject", "")
        if not subject.lower().startswith("re:"):
            subject = f"Re: {subject}"

        quote_lines = [
            f"> From: {original.get('from', '')}",
            f"> To: {original.get('to', '')}",
            f"> Subject: {original.get('subject', '')}",
            f"> Date: {original.get('date', '')}",
            ">",
        ]
        for line in (original.get("plain") or original.get("html") or "").splitlines():
            quote_lines.append(f"> {line}")
        quoted = "\n".join(quote_lines)

        full_body = f"{body}\n\n{quoted}" if quoted.strip("> ") else body

        msg = EmailMessage()
        msg["From"] = formataddr((self.email.split("@")[0], self.email))
        msg["To"] = to_addr
        msg["Subject"] = subject

        orig_msg_id = original.get("message_id", "").strip()
        if orig_msg_id:
            msg["In-Reply-To"] = orig_msg_id
            refs = original.get("references", "").strip()
            msg["References"] = f"{refs} {orig_msg_id}".strip() if refs else orig_msg_id

        content_type = "text/html" if is_html else "text/plain"
        msg.set_content(full_body, subtype=content_type.split("/")[1])

        with smtplib.SMTP_SSL(self.host, self.port) as server:
            server.login(self.email, self.password)
            server.sendmail(self.email, [to_addr], msg.as_string())

    def send_forward(
        self,
        raw_original: bytes,
        to_addrs: list[str],
        subject: str | None = None,
        body: str | None = None,
        cc_addrs: list[str] | None = None,
        bcc_addrs: list[str] | None = None,
        is_html: bool = False,
    ) -> None:
        """Forward an existing message as a message/rfc822 attachment."""
        from email import message_from_bytes

        orig_msg = message_from_bytes(raw_original)
        orig_subject = decode_header_value(orig_msg.get("Subject"))
        if subject is None:
            subject = orig_subject
        if not subject.lower().startswith("fwd:"):
            subject = f"Fwd: {subject}"

        msg = EmailMessage()
        msg["From"] = formataddr((self.email.split("@")[0], self.email))
        msg["To"] = ", ".join(to_addrs)
        if cc_addrs:
            msg["Cc"] = ", ".join(cc_addrs)
        msg["Subject"] = subject

        if body:
            content_type = "text/html" if is_html else "text/plain"
            msg.set_content(body, subtype=content_type.split("/")[1])
        else:
            msg.set_content("")

        msg.add_attachment(
            orig_msg,
            maintype="message",
            subtype="rfc822",
        )

        all_recipients = list(to_addrs)
        if cc_addrs:
            all_recipients.extend(cc_addrs)
        if bcc_addrs:
            all_recipients.extend(bcc_addrs)

        with smtplib.SMTP_SSL(self.host, self.port) as server:
            server.login(self.email, self.password)
            server.sendmail(self.email, all_recipients, msg.as_string())