215 lines
7.7 KiB
Python
215 lines
7.7 KiB
Python
import imaplib
|
||
import email
|
||
import logging
|
||
import re
|
||
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 = "",
|
||
folder: str = "INBOX"):
|
||
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
|
||
self.folder = folder
|
||
|
||
|
||
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, select_inbox: bool = True):
|
||
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)
|
||
if select_inbox:
|
||
_select_mailbox(conn)
|
||
logger.info("已选择 INBOX")
|
||
return conn
|
||
|
||
|
||
def fetch_unseen_emails(account: EmailAccount, skip_uids: set[bytes] | None = None) -> list[Email]:
|
||
conn = _login_and_prepare(account, select_inbox=False)
|
||
if skip_uids is None:
|
||
skip_uids = set()
|
||
|
||
# 列出所有文件夹，跳过 NoSelect 文件夹
|
||
_, folder_data = conn.list()
|
||
folders = []
|
||
if folder_data:
|
||
for item in folder_data:
|
||
if item is None:
|
||
continue
|
||
line = item.decode("utf-8", errors="replace")
|
||
# 跳过 \NoSelect 文件夹（父文件夹，不能直接 SELECT）
|
||
if "\\NoSelect" in line:
|
||
continue
|
||
m = re.search(r'"([^"]*)"\s*$', line)
|
||
if m:
|
||
folders.append(m.group(1))
|
||
if not folders:
|
||
folders = ["INBOX"]
|
||
logger.info("扫描文件夹: %s", ", ".join(folders))
|
||
|
||
emails = []
|
||
skipped = 0
|
||
for folder in folders:
|
||
try:
|
||
status, _ = conn.select(folder, readonly=True)
|
||
if status != "OK":
|
||
logger.debug("无法选择文件夹: %s", folder)
|
||
continue
|
||
_, data = conn.uid("SEARCH", None, "UNSEEN")
|
||
uids = data[0].split() if data[0] else []
|
||
if not uids:
|
||
continue
|
||
|
||
# 先过滤掉正在处理的 UID，避免下载不需要的邮件
|
||
new_uids = [uid for uid in uids if uid not in skip_uids]
|
||
skipped += len(uids) - len(new_uids)
|
||
if not new_uids:
|
||
continue
|
||
logger.info(" %s: %d 封未读 (%d 封跳过)", folder, len(new_uids), len(uids) - len(new_uids))
|
||
|
||
for uid in new_uids:
|
||
try:
|
||
_, msg_data = conn.uid("FETCH", uid, "RFC822")
|
||
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,
|
||
folder=folder))
|
||
|
||
# FETCH 可能被服务器自动标已读，立即标回未读
|
||
conn.uid("STORE", uid, "-FLAGS", "\\Seen")
|
||
except Exception as e:
|
||
logger.error(" UID %s FETCH 失败，跳过: %s", uid, e)
|
||
continue
|
||
except Exception as e:
|
||
logger.warning("扫描文件夹 %s 失败: %s", folder, e)
|
||
continue
|
||
|
||
try:
|
||
conn.logout()
|
||
except Exception:
|
||
pass
|
||
logger.info("共获取 %d 封新邮件", len(emails))
|
||
return emails
|
||
|
||
|
||
@retry()
|
||
def mark_as_seen(account: EmailAccount, uids_with_folder: list[tuple[bytes, str]]):
|
||
if not uids_with_folder:
|
||
return
|
||
logger.info("标记 %d 封邮件为已读", len(uids_with_folder))
|
||
conn = _login_and_prepare(account, select_inbox=False)
|
||
# 按文件夹分组
|
||
from collections import defaultdict
|
||
by_folder: dict[str, list[bytes]] = defaultdict(list)
|
||
for uid, folder in uids_with_folder:
|
||
by_folder[folder].append(uid)
|
||
for folder, uids in by_folder.items():
|
||
try:
|
||
conn.select(folder, readonly=False)
|
||
for uid in uids:
|
||
conn.uid("STORE", uid, "+FLAGS", "\\Seen")
|
||
logger.info(" %s: 标记 %d 封", folder, len(uids))
|
||
except Exception as e:
|
||
logger.error(" 标记 %s 失败: %s", folder, e)
|
||
conn.logout()
|
||
logger.info("标记完成")
|