ai-mail-bot/src/email_client.py

import imaplib
import email
import logging
from email.header import decode_header
from email.utils import parsedate_to_datetime, parseaddr
from typing import Optional
from src.config import EmailAccount
from src.retry import retry

logger = logging.getLogger(__name__)


class Email:
    def __init__(self, uid: bytes, subject: str, sender: str, recipient: str,
                 body: str, date: str, reply_to: str = "", account_email: str = ""):
        self.uid = uid
        self.subject = subject
        self.sender = sender
        self.recipient = recipient
        self.body = body
        self.date = date
        self.reply_to = reply_to
        self.account_email = account_email


def _decode_str(s: str) -> str:
    parts = decode_header(s)
    result = []
    for part, charset in parts:
        if isinstance(part, bytes):
            try:
                result.append(part.decode(charset or "utf-8", errors="replace"))
            except LookupError:
                result.append(part.decode("utf-8", errors="replace"))
        else:
            result.append(part)
    return "".join(result)


def _decode_address_header(header_value: str) -> str:
    if not header_value:
        return ""
    decoded = _decode_str(header_value)
    name, addr = parseaddr(decoded)
    if not addr:
        return name
    if name:
        return f"{name} <{addr}>"
    return addr


def _get_text_from_msg(msg) -> str:
    if msg.is_multipart():
        for part in msg.walk():
            content_type = part.get_content_type()
            if content_type == "text/plain":
                payload = part.get_payload(decode=True)
                if payload:
                    return payload.decode("utf-8", errors="replace")
            elif content_type == "text/html":
                payload = part.get_payload(decode=True)
                if payload:
                    return payload.decode("utf-8", errors="replace")
    else:
        payload = msg.get_payload(decode=True)
        if payload:
            return payload.decode("utf-8", errors="replace")
    return ""


_PROVIDERS_NEED_ID = ("163.com", "126.com")


def _check_provider(email_addr: str, providers: tuple[str, ...]) -> bool:
    return any(email_addr.endswith(p) for p in providers)


def _send_id_command(conn):
    imaplib.Commands["ID"] = "AUTH"
    args = ("name", "AI-Mail-Bot", "version", "1.0.0", "vendor", "personal")
    conn._simple_command("ID", '("' + '" "'.join(args) + '")')


def _select_mailbox(conn, mailbox: str = "INBOX"):
    status, data = conn.select(mailbox)
    if status != "OK":
        raise RuntimeError(f"无法选择邮箱 {mailbox}: {data}")


def _login_and_prepare(account: EmailAccount, timeout: int = 30):
    logger.info("连接 %s:%d (超时%ds)", account.imap_server, account.imap_port, timeout)
    if account.use_ssl:
        conn = imaplib.IMAP4_SSL(account.imap_server, account.imap_port, timeout=timeout)
    else:
        conn = imaplib.IMAP4(account.imap_server, account.imap_port, timeout=timeout)
        if account.use_starttls:
            conn.starttls()
    conn.login(account.username, account.password)
    logger.info("登录成功: %s", account.username)
    if _check_provider(account.username, _PROVIDERS_NEED_ID):
        logger.info("发送 ID 命令 (126/163 兼容)")
        _send_id_command(conn)
    _select_mailbox(conn)
    logger.info("已选择 INBOX")
    return conn


def fetch_unseen_emails(account: EmailAccount) -> list[Email]:
    conn = _login_and_prepare(account)

    _, data = conn.uid("SEARCH", None, "UNSEEN")
    uids = data[0].split() if data[0] else []
    logger.info("UNSEEN 数量: %d", len(uids))
    emails = []

    for uid in uids:
        try:
            _, size_data = conn.uid("FETCH", uid, "(RFC822.SIZE)")
            size = 0
            if size_data[0]:
                m = re.search(rb"RFC822\.SIZE (\d+)", size_data[0])
                if m:
                    size = int(m.group(1))
            large = size > 200_000
            if large:
                logger.info("UID %s 大邮件 (%d bytes)，用 PEEK 拉取", uid, size)

            fetch_cmd = "BODY.PEEK[]" if large else "RFC822"
            _, msg_data = conn.uid("FETCH", uid, fetch_cmd)
            if msg_data[0] is None:
                logger.warning("UID %s FETCH 返回空，跳过", uid)
                continue
            raw_email = msg_data[0][1]
            msg = email.message_from_bytes(raw_email)

            subject = _decode_str(msg.get("Subject", ""))
            sender = _decode_address_header(msg.get("From", ""))
            recipient = _decode_address_header(msg.get("To", ""))
            reply_to = _decode_address_header(msg.get("Reply-To", ""))
            date_str = msg.get("Date", "")
            body = _get_text_from_msg(msg)
            body_len = len(body)

            logger.info("  邮件: [%s] from=%s to=%s reply-to=%s (%d 字符)", subject, sender, recipient, reply_to, body_len)
            emails.append(Email(uid=uid, subject=subject, sender=sender, recipient=recipient,
                                body=body, date=date_str, reply_to=reply_to, account_email=account.username))
        except Exception as e:
            logger.error("UID %s FETCH 失败，跳过并重建连接: %s", uid, e)
            try:
                conn.logout()
            except Exception:
                pass
            try:
                conn = _login_and_prepare(account)
            except Exception as e2:
                logger.error("重建连接失败: %s", e2)
                break
            continue

    try:
        conn.logout()
    except Exception:
        pass
    logger.info("共获取 %d 封新邮件", len(emails))
    return emails


@retry()
def mark_as_seen(account: EmailAccount, uids: list[bytes]):
    if not uids:
        return
    logger.info("标记 %d 封邮件为已读", len(uids))
    conn = _login_and_prepare(account)
    for uid in uids:
        conn.uid("STORE", uid, "+FLAGS", "\\Seen")
    conn.logout()
    logger.info("标记完成")