group_fqcd_jr/app/service/offsite_mail_adapter.py

"""场外基金真实邮件接收和原始文件保存适配层。"""

from __future__ import annotations

import hashlib
import imaplib
import os
import re
from contextlib import suppress
from dataclasses import dataclass
from datetime import UTC, datetime
from email import message_from_bytes
from email.header import decode_header, make_header
from email.message import EmailMessage, Message
from email.utils import parseaddr
from pathlib import Path
from typing import Protocol, cast

from app.core.config import Settings


@dataclass(frozen=True)
class RawMailAttachment:
    filename: str
    media_type: str
    payload: bytes


@dataclass(frozen=True)
class RawMailMessage:
    imap_uid: str
    message_id: str
    sender: str
    return_path: str | None
    auth_result: dict[str, object]
    raw_message: bytes
    received_at: datetime
    attachments: tuple[RawMailAttachment, ...]


@dataclass(frozen=True)
class SavedMailAttachment:
    filename: str
    file_hash: str
    media_type: str
    size_bytes: int
    original_file_path: str


@dataclass(frozen=True)
class SavedMailMessage:
    imap_uid: str
    message_id: str
    sender: str
    return_path: str | None
    auth_result: dict[str, object]
    eml_path: str
    attachments: tuple[SavedMailAttachment, ...]


class ImapConnection(Protocol):
    def login(self, user: str, password: str) -> tuple[str, list[bytes]]: ...

    def select(self, mailbox: str, readonly: bool = False) -> tuple[str, list[bytes]]: ...

    def uid(self, command: str, *args: object) -> tuple[str, list[object]]: ...

    def noop(self) -> tuple[str, list[bytes]]: ...

    def send(self, data: bytes) -> object: ...

    def logout(self) -> tuple[str, list[bytes]]: ...


class OffsiteMailStorage:
    """保存原始 eml 和附件，避免覆盖已有原始文件。"""

    def __init__(self, root_dir: str | Path) -> None:
        self.root_dir = Path(root_dir)

    def save(self, mail: RawMailMessage) -> SavedMailMessage:
        folder = self._folder_for(mail)
        folder.mkdir(parents=True, exist_ok=True)
        eml_path = folder / "message.eml"
        self._write_once(eml_path, mail.raw_message)
        saved = []
        for index, attachment in enumerate(mail.attachments, start=1):
            file_hash = hashlib.sha256(attachment.payload).hexdigest()
            filename = f"A{index:02d}_{file_hash[:12]}_{self._safe_name(attachment.filename)}"
            path = folder / filename
            self._write_once(path, attachment.payload)
            saved.append(SavedMailAttachment(
                filename=attachment.filename,
                file_hash=file_hash,
                media_type=attachment.media_type,
                size_bytes=len(attachment.payload),
                original_file_path=str(path),
            ))
        return SavedMailMessage(
            imap_uid=mail.imap_uid,
            message_id=mail.message_id,
            sender=mail.sender,
            return_path=mail.return_path,
            auth_result=mail.auth_result,
            eml_path=str(eml_path),
            attachments=tuple(saved),
        )

    def _folder_for(self, mail: RawMailMessage) -> Path:
        day = mail.received_at.strftime("%Y%m%d")
        source = mail.message_id or mail.imap_uid
        return self.root_dir / day / self._safe_name(source.strip("<>") or mail.imap_uid)

    @staticmethod
    def _safe_name(value: str) -> str:
        safe = re.sub(r"[^0-9A-Za-z._-]+", "_", value.strip())
        return safe[:120] or "unknown"

    @staticmethod
    def _write_once(path: Path, payload: bytes) -> None:
        if path.exists():
            return
        path.write_bytes(payload)
        with suppress(OSError):
            os.chmod(path, 0o444)


class OffsiteImapReceiver:
    """只监听收件箱，按 UID 增量拉取原始邮件。"""

    inbox_name = "INBOX"

    def __init__(self, settings: Settings, connection: ImapConnection | None = None) -> None:
        self.settings = settings
        self._connection = connection
        self.last_scanned_uid: str | None = None

    def health_check(self) -> dict[str, object]:
        if not self.settings.offsite_imap_enabled:
            return {"status": "disabled", "message": "场外 IMAP 未启用"}
        missing = self._missing_config()
        if missing:
            return {"status": "misconfigured", "missing": missing}
        try:
            connection = self._ensure_connection()
            connection.noop()
            return {"status": "ok", "mailbox": self.settings.offsite_mailbox}
        except Exception as exc:
            self.close()
            return {"status": "error", "message": type(exc).__name__}

    def fetch_since(self, last_uid: str | None, *, limit: int = 20) -> tuple[RawMailMessage, ...]:
        if not self.settings.offsite_imap_enabled:
            return ()
        self.last_scanned_uid = None
        missing = self._missing_config()
        if missing:
            raise RuntimeError(f"场外IMAP配置缺失：{', '.join(missing)}")
        connection = self._ensure_connection()
        self._select_inbox(connection)
        start_uid = int(last_uid) + 1 if last_uid and last_uid.isdigit() else 1
        status, data = connection.uid("search", None, f"UID {start_uid}:*")
        search_payload = self._first_bytes(data)
        if status != "OK" or not search_payload:
            return ()
        uid_values = search_payload.split()[:limit]
        if uid_values:
            self.last_scanned_uid = uid_values[-1].decode("ascii")
        messages: list[RawMailMessage] = []
        for uid in uid_values:
            item = self._fetch_one(connection, uid.decode("ascii"))
            if item and item.sender in self.settings.offsite_allowed_senders:
                messages.append(item)
        return tuple(messages)

    def wait_for_new_mail(self, timeout_seconds: float) -> bool:
        """使用 IMAP IDLE 等待新邮件，超时返回 False。"""
        if not self.settings.offsite_imap_enabled or not self.settings.offsite_imap_idle_enabled:
            return False
        missing = self._missing_config()
        if missing:
            raise RuntimeError(f"场外IMAP配置缺失：{', '.join(missing)}")
        connection = self._ensure_connection()
        idle_tag: object | None = None
        try:
            self._select_inbox(connection)
            idle_tag = self._idle_start(connection)
            return self._idle_wait(connection, timeout_seconds)
        except Exception:
            self.close()
            raise
        finally:
            if idle_tag is not None and self._connection is connection:
                try:
                    self._idle_done(connection, idle_tag, timeout_seconds)
                except Exception:
                    self.close()
                    raise

    def close(self) -> None:
        if self._connection is None:
            return
        try:
            self._connection.logout()
        finally:
            self._connection = None

    def _ensure_connection(self) -> ImapConnection:
        if self._connection is not None:
            return self._connection
        client_cls = imaplib.IMAP4_SSL if self.settings.offsite_imap_use_ssl else imaplib.IMAP4
        connection = cast(
            ImapConnection,
            client_cls(self.settings.offsite_imap_host, self.settings.offsite_imap_port),
        )
        connection.login(self.settings.offsite_imap_username, self.settings.offsite_imap_password)
        self._send_client_id(connection)
        self._connection = connection
        return connection

    @staticmethod
    def _send_client_id(connection: ImapConnection) -> None:
        """向支持 RFC 2971 的邮箱声明客户端身份，兼容不支持该命令的服务。"""
        simple_command = getattr(connection, "_simple_command", None)
        if not callable(simple_command):
            return
        imaplib.Commands.setdefault("ID", ("NONAUTH", "AUTH", "SELECTED"))
        simple_command(
            "ID",
            '("name" "NanfangFund" "version" "1.0" "vendor" "internal")',
        )

    def _select_inbox(self, connection: ImapConnection) -> None:
        status, _ = connection.select(self.inbox_name, readonly=True)
        if status != "OK":
            raise RuntimeError(f"场外 IMAP 选择 {self.inbox_name} 失败：{status}")

    @staticmethod
    def _idle_start(connection: ImapConnection) -> object:
        custom_start = getattr(connection, "idle_start", None)
        if callable(custom_start):
            return custom_start()
        new_tag = getattr(connection, "_new_tag", None)
        get_response = getattr(connection, "_get_response", None)
        if not callable(new_tag) or not callable(get_response):
            raise RuntimeError("当前 IMAP 连接不支持 IDLE")
        tag = new_tag()
        connection.send(tag + b" IDLE\r\n")
        response = get_response()
        if response is not None:
            raise imaplib.IMAP4.error("IMAP IDLE 未收到继续响应")
        return tag

    @staticmethod
    def _idle_wait(connection: ImapConnection, timeout_seconds: float) -> bool:
        custom_wait = getattr(connection, "idle_wait", None)
        if callable(custom_wait):
            return bool(custom_wait(timeout_seconds))
        get_response = getattr(connection, "_get_response", None)
        sock = getattr(connection, "sock", None)
        if not callable(get_response) or sock is None:
            raise RuntimeError("当前 IMAP 连接不支持 IDLE 等待")
        previous_timeout = sock.gettimeout()
        sock.settimeout(timeout_seconds)
        try:
            response = get_response()
        except TimeoutError:
            return False
        finally:
            sock.settimeout(previous_timeout)
        if not isinstance(response, bytes):
            return False
        upper_response = response.upper()
        return any(marker in upper_response for marker in (b" EXISTS", b" RECENT"))

    @staticmethod
    def _idle_done(
        connection: ImapConnection, idle_tag: object, timeout_seconds: float
    ) -> None:
        custom_done = getattr(connection, "idle_done", None)
        if callable(custom_done):
            custom_done()
            return
        get_tagged_response = getattr(connection, "_get_tagged_response", None)
        sock = getattr(connection, "sock", None)
        if not callable(get_tagged_response) or sock is None:
            raise RuntimeError("当前 IMAP 连接不支持退出 IDLE")
        previous_timeout = sock.gettimeout()
        sock.settimeout(timeout_seconds)
        try:
            connection.send(b"DONE\r\n")
            response_type, _ = get_tagged_response(idle_tag)
            if response_type != "OK":
                raise imaplib.IMAP4.error(f"退出 IMAP IDLE 失败：{response_type}")
        finally:
            sock.settimeout(previous_timeout)

    def _fetch_one(self, connection: ImapConnection, uid: str) -> RawMailMessage | None:
        status, data = connection.uid("fetch", uid, "(RFC822)")
        if status != "OK" or not data:
            return None
        raw = self._extract_rfc822(data)
        if raw is None:
            return None
        parsed = message_from_bytes(raw)
        message = cast(EmailMessage, parsed)
        sender = parseaddr(message.get("From", ""))[1]
        return_path = parseaddr(message.get("Return-Path", ""))[1] or None
        message_id = message.get("Message-ID", f"<imap-{uid}>")
        return RawMailMessage(
            imap_uid=uid,
            message_id=message_id,
            sender=sender,
            return_path=return_path,
            auth_result=self._auth_result(message),
            raw_message=raw,
            received_at=datetime.now(UTC).replace(tzinfo=None),
            attachments=self._attachments(message),
        )

    @staticmethod
    def _first_bytes(data: list[object]) -> bytes | None:
        for item in data:
            if isinstance(item, bytes):
                return item
        return None

    @staticmethod
    def _extract_rfc822(data: list[object]) -> bytes | None:
        for item in data:
            if isinstance(item, tuple):
                payload = item[1]
                if isinstance(payload, bytes):
                    return payload
            if isinstance(item, bytes) and item.startswith(b"From:"):
                return item
        return None

    @staticmethod
    def _attachments(message: Message) -> tuple[RawMailAttachment, ...]:
        attachments: list[RawMailAttachment] = []
        for part in message.walk():
            if part.is_multipart():
                continue
            filename = part.get_filename()
            disposition = part.get_content_disposition()
            if not filename and disposition != "attachment":
                continue
            raw_payload = part.get_payload(decode=True)
            payload = raw_payload if isinstance(raw_payload, bytes) else b""
            attachments.append(RawMailAttachment(
                filename=OffsiteImapReceiver._decode_filename(filename) or "attachment.bin",
                media_type=part.get_content_type(),
                payload=payload,
            ))
        return tuple(attachments)

    @staticmethod
    def _decode_filename(filename: str | None) -> str:
        if not filename:
            return ""
        try:
            return str(make_header(decode_header(filename)))
        except (LookupError, UnicodeError, ValueError):
            return filename

    @staticmethod
    def _auth_result(message: Message) -> dict[str, object]:
        return {
            "authentication_results": message.get_all("Authentication-Results", []),
            "received_spf": message.get_all("Received-SPF", []),
            "dkim_signature_present": bool(message.get("DKIM-Signature")),
        }

    def _missing_config(self) -> list[str]:
        missing = []
        if not self.settings.offsite_imap_host:
            missing.append("OFFSITE_IMAP_HOST")
        if not self.settings.offsite_imap_username:
            missing.append("OFFSITE_IMAP_USERNAME")
        if not self.settings.offsite_imap_password:
            missing.append("OFFSITE_IMAP_PASSWORD")
        return missing