services/register/pool_mail.py

"""Account-pool mail backend for the registration engine.

The pool web service (rg-gpt ``pool/app.py``, e.g. ``http://1.94.147.46:8091``)
hands out **real** Outlook / Gmail mailboxes with an atomic lease and records the
outcome so a mailbox that already backed one ChatGPT signup is never re-used.

Unlike the temp-mail HTTP providers, the pool gives credentials rather than an
OTP-fetch API, so the OTP is read over IMAP:

  - Outlook / Hotmail / Live (personal MSA): basic-auth IMAP and ROPC are both
    disabled since Sept 2024, so we mint an OAuth2 access token from the
    long-lived ``refresh_token`` (device-code bootstrapped in rg-gpt) and log in
    with IMAP **XOAUTH2**.
  - Gmail: plain IMAP LOGIN with a 2FA **app password**.

A base mailbox is shared by its ``+1..+5`` plus-address aliases, so messages are
filtered by exact ``To``/``Cc`` recipient to stay concurrency-safe.

This module is deliberately free of app-wide settings / on-disk pool logic (that
lives in the caller); it only holds the pool HTTP client, the MSA token helper,
the IMAP readers, and MIME parsing.
"""
from __future__ import annotations

import base64
import email as email_pkg
import html as html_lib
import imaplib
import re
import time
from datetime import datetime, timezone
from email.header import decode_header, make_header
from email.utils import getaddresses, parsedate_to_datetime
from typing import Any

import requests

# Thunderbird's registered public client — personal MSA accounts accept it for
# the IMAP scope. No client secret (public client).
THUNDERBIRD_CLIENT_ID = "9e5f94bc-e8a4-4e73-b8be-63364c29d753"
# Personal accounts live under /consumers; /common and /organizations reject them.
TOKEN_ENDPOINT = "https://login.microsoftonline.com/consumers/oauth2/v2.0/token"
IMAP_SCOPE = "https://outlook.office.com/IMAP.AccessAsUser.All offline_access"
IMAP_HOST = "outlook.office365.com"
IMAP_PORT = 993
# Folders an OpenAI OTP mail can land in on a fresh mailbox.
IMAP_FOLDERS = ("INBOX", "Junk")

GMAIL_IMAP_HOST = "imap.gmail.com"
GMAIL_FOLDERS = ("INBOX", "[Gmail]/Spam")
GMAIL_DOMAINS = ("gmail.com", "googlemail.com")

ACCESS_TOKEN_SKEW_SECONDS = 60  # refresh a bit before actual expiry
# imaplib.IMAP4_SSL has NO default timeout: without this a blocked TCP connect (e.g. an
# egress that can't reach office365:993) hangs the register thread forever ("卡住"). Bound it
# so a dead route fails fast and the OTP wait loop can retry / time out cleanly.
IMAP_CONNECT_TIMEOUT = 20


class OutlookAuthError(RuntimeError):
    """Raised when an MSA token refresh or an IMAP login cannot complete."""


# ---- pool HTTP client ---------------------------------------------------------


class PoolClient:
    """Thin client for the pool web service (rg-gpt pool/app.py)."""

    def __init__(self, base_url: str, api_key: str, timeout: int = 30) -> None:
        self.base = str(base_url or "").rstrip("/")
        self.key = str(api_key or "")
        self.timeout = timeout
        self._s = requests.Session()
        self._s.trust_env = False  # never route the pool call through an OpenAI proxy

    def _headers(self) -> dict[str, str]:
        return {"X-API-Key": self.key, "Content-Type": "application/json"}

    def lease(self, count: int = 1, leased_by: str = "", kind: str = "") -> list[dict[str, Any]]:
        resp = self._s.post(
            f"{self.base}/api/lease", headers=self._headers(),
            json={"count": count, "leased_by": leased_by, "kind": kind}, timeout=self.timeout,
        )
        resp.raise_for_status()
        return resp.json().get("leased") or []

    def available(self, kind: str = "") -> int:
        resp = self._s.get(
            f"{self.base}/api/available", headers=self._headers(),
            params={"kind": kind}, timeout=self.timeout,
        )
        resp.raise_for_status()
        return int(resp.json().get("available") or 0)

    def report(self, acct_id: int, status: str, **fields: Any) -> dict[str, Any]:
        resp = self._s.post(
            f"{self.base}/api/accounts/{acct_id}/result", headers=self._headers(),
            json={"status": status, **fields}, timeout=self.timeout,
        )
        resp.raise_for_status()
        return resp.json()

    def fetch_otp(self, acct_id: int, lease_token: str = "", since: float = 0.0, exclude: str = "") -> str | None:
        """Ask the pool to read the mailbox server-side (its China egress can reach office365,
        unlike this consumer) and return the newest matching OpenAI/ChatGPT OTP, or None."""
        resp = self._s.get(
            f"{self.base}/api/accounts/{acct_id}/otp", headers=self._headers(),
            params={"lease_token": lease_token, "since": since, "exclude": exclude}, timeout=self.timeout,
        )
        resp.raise_for_status()
        code = resp.json().get("code")
        return str(code) if code else None


# ---- MSA token ----------------------------------------------------------------


def refresh_access_token(refresh_token: str, *, client_id: str = THUNDERBIRD_CLIENT_ID) -> dict[str, Any]:
    resp = requests.post(
        TOKEN_ENDPOINT,
        data={"grant_type": "refresh_token", "client_id": client_id,
              "refresh_token": refresh_token, "scope": IMAP_SCOPE},
        timeout=30,
    )
    data = resp.json()
    if "access_token" not in data:
        raise OutlookAuthError(f"refresh failed: {data.get('error')} {data.get('error_description')}")
    return data


# ---- IMAP readers -------------------------------------------------------------


def _xoauth2_bytes(email: str, access_token: str) -> bytes:
    return f"user={email}\x01auth=Bearer {access_token}\x01\x01".encode()


def _imap_collect(conn: imaplib.IMAP4_SSL, folders: tuple[str, ...], per_folder: int) -> list[dict[str, Any]]:
    """After the connection is authenticated, collect recent messages (newest first)."""
    out: list[dict[str, Any]] = []
    for folder in folders:
        typ, _ = conn.select(folder, readonly=True)
        if typ != "OK":
            continue
        typ, data = conn.search(None, "ALL")
        if typ != "OK" or not data or not data[0]:
            continue
        ids = data[0].split()
        for msg_id in reversed(ids[-per_folder:]):
            typ, msg_data = conn.fetch(msg_id, "(RFC822)")
            if typ != "OK" or not msg_data or not msg_data[0]:
                continue
            msg = email_pkg.message_from_bytes(msg_data[0][1])
            subject, text_body, html_body, from_addr = _parse_message(msg)
            out.append({
                "folder": folder, "subject": subject, "text": text_body, "html": html_body,
                "sender": from_addr, "recipients": _msg_recipients(msg),
                "received_at": _parse_dt(msg.get("Date")),
            })
    out.sort(key=lambda m: m["received_at"] or datetime.min.replace(tzinfo=timezone.utc), reverse=True)
    return out


def imap_read_messages(email: str, access_token: str, *, per_folder: int = 20) -> list[dict[str, Any]]:
    """Outlook: recent messages across INBOX + Junk via XOAUTH2."""
    conn = imaplib.IMAP4_SSL(IMAP_HOST, IMAP_PORT, timeout=IMAP_CONNECT_TIMEOUT)
    try:
        try:
            conn.authenticate("XOAUTH2", lambda _challenge: _xoauth2_bytes(email, access_token))
        except imaplib.IMAP4.error as exc:
            raise OutlookAuthError(f"XOAUTH2 login failed for {email}: {exc}") from exc
        return _imap_collect(conn, IMAP_FOLDERS, per_folder)
    finally:
        try:
            conn.logout()
        except Exception:
            pass


def gmail_imap_read_messages(login_email: str, app_password: str, *, per_folder: int = 20) -> list[dict[str, Any]]:
    """Gmail: recent messages via plain IMAP LOGIN with an app password (2FA + app pw)."""
    conn = imaplib.IMAP4_SSL(GMAIL_IMAP_HOST, IMAP_PORT, timeout=IMAP_CONNECT_TIMEOUT)
    try:
        try:
            conn.login(login_email, app_password)
        except imaplib.IMAP4.error as exc:
            raise OutlookAuthError(f"Gmail IMAP login failed for {login_email}: {exc}") from exc
        return _imap_collect(conn, GMAIL_FOLDERS, per_folder)
    finally:
        try:
            conn.logout()
        except Exception:
            pass


# ---- MIME parsing -------------------------------------------------------------


def _parse_dt(value: Any) -> datetime | None:
    text = str(value or "").strip()
    if not text:
        return None
    try:
        dt = parsedate_to_datetime(text)
        if dt is None:
            return None
        return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)
    except Exception:
        return None


def _decode_mime_header(value: str | None) -> str:
    if not value:
        return ""
    try:
        return str(make_header(decode_header(value)))
    except Exception:
        return str(value)


def _part_to_text(part: Any) -> str:
    try:
        payload = part.get_payload(decode=True)
        if payload is None:
            return str(part.get_payload())
        charset = part.get_content_charset() or "utf-8"
        return payload.decode(charset, errors="replace")
    except Exception:
        try:
            return str(part.get_payload())
        except Exception:
            return ""


def _parse_message(msg: Any) -> tuple[str, str, str, str]:
    subject = _decode_mime_header(msg.get("Subject", ""))
    from_addr = _decode_mime_header(msg.get("From", ""))
    text_body = ""
    html_body = ""
    if msg.is_multipart():
        for part in msg.walk():
            if part.is_multipart():
                continue
            ctype = part.get_content_type()
            if "attachment" in (part.get("Content-Disposition") or "").lower():
                continue
            if ctype == "text/plain" and not text_body:
                text_body = _part_to_text(part)
            elif ctype == "text/html" and not html_body:
                html_body = _part_to_text(part)
    else:
        decoded = _part_to_text(msg)
        if msg.get_content_type() == "text/html":
            html_body = decoded
        else:
            text_body = decoded
    return subject, text_body, html_body, from_addr


def _msg_recipients(msg: Any) -> set[str]:
    """Lowercased To+Cc addresses — these preserve the +tag for plus-addressed mail
    (Delivered-To is always the base mailbox, so it is deliberately excluded)."""
    vals: list[str] = []
    for header in ("To", "Cc"):
        vals.extend(msg.get_all(header, []) or [])
    return {addr.strip().lower() for _, addr in getaddresses(vals) if addr and "@" in addr}


def base_email(email: str) -> str:
    """Strip a +tag sub-address: user+2@outlook.com -> user@outlook.com."""
    local, sep, domain = str(email or "").strip().partition("@")
    if not sep:
        return str(email or "").strip().lower()
    return f"{local.split('+', 1)[0]}@{domain}".lower()


def html_to_text(value: Any) -> str:
    content = str(value or "")
    if not content:
        return ""
    content = re.sub(r"(?is)<(script|style)\b.*?>.*?</\1>", " ", content)
    content = re.sub(r"(?is)<!--.*?-->", " ", content)
    content = re.sub(r"(?i)<br\s*/?>", "\n", content)
    content = re.sub(r"(?i)</(?:p|div|tr|table|h[1-6]|li|td|section|article)>", "\n", content)
    content = re.sub(r"(?s)<[^>]+>", " ", content)
    content = html_lib.unescape(content)
    content = re.sub(r"[\t\r\f\v ]+", " ", content)
    content = re.sub(r"\n\s+", "\n", content)
    return content.strip()