webmail-client/backend/app/providers/imap_smtp/provider.py

"""``ImapSmtpProvider`` — IMAP (aioimaplib) + SMTP (aiosmtplib) MailProvider.

One instance == one IMAP connection (SPEC §6.3 caps connections per account);
the sync engine creates separate instances for the IDLE loop vs. backfill/jobs.
Credentials arrive already decrypted (decryption lives only in workers /
accounts.test per SPEC §10.2).
"""

from __future__ import annotations

import asyncio
import contextlib
import re
import ssl
from collections.abc import AsyncIterator
from dataclasses import dataclass
from datetime import date

import aiosmtplib
from aioimaplib import aioimaplib

from app.concurrency import loop_semaphore
from app.config import get_settings, smtp_helo_name
from app.providers.base import (
    FlagChange,
    MailboxDelta,
    MailProvider,
    ProviderAuthError,
    ProviderConnectError,
    ProviderError,
    ProviderNotSupported,
    ProviderRejectedError,
    RemoteMailbox,
    RemoteMessage,
)
from app.providers.imap_smtp.parsing import (
    encode_mailbox_name,
    parse_fetch_lines,
    parse_list_line,
    parse_search_uids,
    parse_select_untagged,
    parse_vanished,
)

_FETCH_ITEMS = "(UID FLAGS INTERNALDATE RFC822.SIZE BODY.PEEK[])"
# Gmail (X-GM-EXT-1): also pull each message's labels so custom labels can map to
# chips without visiting every label folder (SPEC: Gmail sync / no-duplication).
_FETCH_ITEMS_GM = "(UID FLAGS INTERNALDATE RFC822.SIZE X-GM-LABELS BODY.PEEK[])"
# SPEC §5.1 asks for IDLE to be re-issued at least every 25 min; 10 also means a
# connection that died silently (no FIN, e.g. the network dropped) is noticed
# within ~11 min — the re-issue finds it — instead of ~29.
_IDLE_TIMEOUT = 10 * 60
_CONNECT_TIMEOUT = 20
_PER_ACCOUNT_CONN_CAP = 4  # SPEC §6.3 / §4: IDLE + transient + short-lived jobs
_APPENDUID_RE = re.compile(rb"APPENDUID\s+\d+\s+(\d+)")
_COPYUID_RE = re.compile(rb"COPYUID\s+\d+\s+\S+\s+(\d+)")
_IMAP_MONTHS = ("Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec")


class _KeepConnectTask:
    """aioimaplib starts the TCP connection as a task it never awaits: a refused
    or unreachable server then surfaces only as a 20 s greeting timeout, plus a
    "Task exception was never retrieved" traceback in the log. Keep the task, so
    ``_connect`` can await it — failing at once, with the real reason."""

    connect_task: asyncio.Task[object]

    def create_client(  # type: ignore[no-untyped-def]
        self, host, port, loop, conn_lost_cb=None, ssl_context=None
    ) -> None:
        # This replaces IMAP4_SSL.create_client too (it comes first in the MRO),
        # so do its one job here: without a context the connection is plain
        # TCP to port 993, and the server just waits for a TLS handshake.
        if ssl_context is None and isinstance(self, aioimaplib.IMAP4_SSL):
            ssl_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
        local_loop = loop if loop is not None else asyncio.get_running_loop()
        self.protocol = aioimaplib.IMAP4ClientProtocol(local_loop, conn_lost_cb)
        self.connect_task = local_loop.create_task(
            local_loop.create_connection(lambda: self.protocol, host, port, ssl=ssl_context)
        )


class _IMAP4(_KeepConnectTask, aioimaplib.IMAP4):  # type: ignore[misc]
    pass


class _IMAP4_SSL(_KeepConnectTask, aioimaplib.IMAP4_SSL):  # type: ignore[misc]  # noqa: N801
    pass


def _smtp_error(exc: aiosmtplib.SMTPException) -> ProviderError:
    """An SMTP failure as the provider error that says whether to try again.

    A 5xx reply is the server's final word (recipient unknown, message refused):
    ``ProviderRejectedError``, or ``ProviderAuthError`` for a refused login. Any
    4xx reply, a timeout or a dropped connection is temporary:
    ``ProviderConnectError``, and the send is retried later.
    """
    if isinstance(exc, aiosmtplib.SMTPRecipientsRefused):
        codes = [r.code for r in exc.recipients]
        if codes and all(c >= 500 for c in codes):
            return ProviderRejectedError(f"recipients refused: {exc}")
        return ProviderConnectError(f"SMTP send failed: {exc}")
    if isinstance(exc, aiosmtplib.SMTPResponseException) and exc.code >= 500:
        if isinstance(exc, aiosmtplib.SMTPAuthenticationError):
            return ProviderAuthError(f"SMTP auth failed: {exc}")
        return ProviderRejectedError(f"SMTP refused: {exc}")
    return ProviderConnectError(f"SMTP send failed: {exc}")


def _imap_date(d: date) -> str:
    """RFC 3501 search date, e.g. ``5-Jul-2026`` (locale-independent)."""
    return f"{d.day}-{_IMAP_MONTHS[d.month - 1]}-{d.year}"


def _date_criteria(since: date | None, before: date | None) -> str:
    parts: list[str] = []
    if since is not None:
        parts.append(f"SINCE {_imap_date(since)}")
    if before is not None:
        parts.append(f"BEFORE {_imap_date(before)}")
    return " ".join(parts) if parts else "ALL"


async def _require_public(host: str) -> None:
    """Refuse a user-named server that now resolves to a non-public address."""
    from app.accounts.ssrf import SsrfError, assert_host_is_public

    try:
        await assert_host_is_public(host)
    except SsrfError as exc:
        raise ProviderConnectError(f"refusing to connect to {host}: {exc}") from exc


@dataclass(frozen=True)
class ProviderConfig:
    imap_host: str
    imap_port: int
    imap_tls: str  # "ssl" | "starttls" | "plain"
    smtp_host: str
    smtp_port: int
    smtp_tls: str  # "ssl" | "starttls" | "plain"
    username: str
    password: str  # plaintext; caller decrypts per §10.2
    # TLS settings other than the default (certificate checked): the hosted
    # mail server on the internal network (app/hosted/native.py).
    tls_context: ssl.SSLContext | None = None
    # The hosted mail server: every user's mailbox is on that one host, so the
    # per-host connection cap (which spares outside providers) does not apply.
    internal: bool = False
    # A server the user named (every external account): its name is checked
    # again before each connection, so DNS changed after the account was added
    # cannot turn the worker on the internal network (DNS rebinding).
    public_only: bool = False


class ImapSmtpProvider(MailProvider):
    def __init__(self, config: ProviderConfig) -> None:
        self._cfg = config
        self._imap: aioimaplib.IMAP4 | None = None
        self._selected: str | None = None
        self._uidvalidity_seen: dict[str, int] = {}
        self._conn_slots: list[asyncio.Semaphore] = []
        self._qresync_enabled = False
        self._gm_ext = False

    @property
    def gmail_labels_supported(self) -> bool:
        """True when the server advertises ``X-GM-EXT-1`` (Gmail), so FETCH can pull
        ``X-GM-LABELS`` and the engine can use the no-duplication Gmail sync path."""
        return self._gm_ext

    @property
    def qresync_enabled(self) -> bool:
        """True once ``ENABLE QRESYNC`` succeeded on this connection (RFC 7162).

        When set, :meth:`poll` reports server-side expunges via VANISHED; otherwise
        deletion detection falls back to the engine's full-UID reconcile."""
        return self._qresync_enabled

    # --- connection lifecycle ------------------------------------------------

    async def _acquire_conn_slots(self) -> None:
        """Bound simultaneous open sockets: ≤4 per account and ≤cap per host (§4).

        Held for the connection's whole lifetime and released in :meth:`close`.
        """
        cfg = self._cfg
        account = loop_semaphore(
            "imap_account", f"{cfg.username}@{cfg.imap_host}", _PER_ACCOUNT_CONN_CAP
        )
        await account.acquire()  # consistent order (account, then host) avoids deadlock
        if cfg.internal:
            self._conn_slots = [account]
            return
        host = loop_semaphore("imap_host", cfg.imap_host, get_settings().per_host_conn_cap)
        try:
            await host.acquire()
        except BaseException:
            account.release()
            raise
        self._conn_slots = [host, account]

    def _release_conn_slots(self) -> None:
        for sem in self._conn_slots:
            sem.release()
        self._conn_slots = []

    async def _connect(self) -> aioimaplib.IMAP4:
        if self._imap is not None:
            return self._imap
        cfg = self._cfg
        if cfg.public_only:
            await _require_public(cfg.imap_host)
        await self._acquire_conn_slots()
        try:
            try:
                imap: aioimaplib.IMAP4
                if cfg.imap_tls == "ssl":
                    imap = _IMAP4_SSL(
                        host=cfg.imap_host,
                        port=cfg.imap_port,
                        timeout=_CONNECT_TIMEOUT,
                        ssl_context=cfg.tls_context,
                    )
                else:
                    imap = _IMAP4(host=cfg.imap_host, port=cfg.imap_port, timeout=_CONNECT_TIMEOUT)
                # Refused / unknown host: fails here at once. No answer at all:
                # gives up after the timeout (the connect attempt is cancelled).
                await asyncio.wait_for(imap.connect_task, _CONNECT_TIMEOUT)
                await imap.wait_hello_from_server()
                if cfg.imap_tls == "starttls":
                    await imap.starttls()
            except (OSError, TimeoutError) as exc:
                why = str(exc) or f"no answer within {_CONNECT_TIMEOUT} s"
                raise ProviderConnectError(f"cannot connect to {cfg.imap_host}: {why}") from exc

            resp = await imap.login(cfg.username, cfg.password)
            if resp.result != "OK":
                detail = (
                    b" ".join(resp.lines).decode("utf-8", "ignore")
                    if resp.lines
                    else "login failed"
                )
                raise ProviderAuthError(detail)
        except BaseException:
            self._release_conn_slots()  # never hold a slot without an open socket
            raise
        self._imap = imap
        caps = {c.upper() for c in (imap.protocol.capabilities or set())}
        self._gm_ext = "X-GM-EXT-1" in caps
        await self._maybe_enable_qresync(imap)
        return imap

    async def _maybe_enable_qresync(self, imap: aioimaplib.IMAP4) -> None:
        """Enable QRESYNC when advertised so SELECT reports VANISHED (RFC 7162).

        Best-effort: servers without QRESYNC (e.g. Gmail, which is CONDSTORE-only)
        simply keep ``_qresync_enabled`` False and use the engine's reconcile path."""
        caps = {c.upper() for c in (imap.protocol.capabilities or set())}
        if "QRESYNC" not in caps:
            return
        try:
            resp = await imap.enable("QRESYNC")
            self._qresync_enabled = resp.result == "OK"
        except Exception:  # noqa: BLE001 - degrade gracefully to the reconcile path
            self._qresync_enabled = False

    @staticmethod
    def _wire(mailbox: str) -> str:
        """Server wire form of a mailbox name: mUTF-7 encoded and quoted. Handles
        non-ASCII names (Gmail's localized ``[Gmail]/Wysłane``) and names with
        spaces (``[Gmail]/All Mail``); ASCII names are unchanged bar the quotes."""
        return str(aioimaplib.quoted(encode_mailbox_name(mailbox)))

    async def _select(self, mailbox: str) -> dict[str, int]:
        imap = await self._connect()
        resp = await imap.select(self._wire(mailbox))
        if resp.result != "OK":
            raise ProviderConnectError(f"cannot select {mailbox!r}")
        self._selected = mailbox
        return parse_select_untagged(resp.lines)

    async def close(self) -> None:
        if self._imap is not None:
            with contextlib.suppress(Exception):  # best-effort teardown
                await self._imap.logout()
            self._imap = None
            self._selected = None
            self._release_conn_slots()  # free the per-account / per-host slots

    # --- MailProvider protocol ----------------------------------------------

    async def verify(self) -> None:
        imap = await self._connect()
        resp = await imap.list('""', "*")
        if resp.result != "OK":
            raise ProviderConnectError("LIST failed after login")

    async def list_mailboxes(self) -> list[RemoteMailbox]:
        imap = await self._connect()
        resp = await imap.list('""', "*")
        if resp.result != "OK":
            raise ProviderConnectError("LIST failed")
        out: list[RemoteMailbox] = []
        for line in resp.lines:
            if not isinstance(line, (bytes, bytearray)):
                continue
            parsed = parse_list_line(bytes(line))
            if parsed is None:
                continue
            uidvalidity = await self._mailbox_uidvalidity(parsed.path)
            out.append(
                RemoteMailbox(
                    path=parsed.path, special_use=parsed.special_use, uidvalidity=uidvalidity
                )
            )
        return out

    async def _mailbox_uidvalidity(self, path: str) -> int:
        imap = await self._connect()
        try:
            resp = await imap.status(self._wire(path), "(UIDVALIDITY)")
        except Exception:
            return 0
        if resp.result != "OK":
            return 0
        for line in resp.lines:
            if isinstance(line, (bytes, bytearray)):
                m = re.search(rb"UIDVALIDITY\s+(\d+)", bytes(line))
                if m:
                    return int(m.group(1))
        return 0

    async def backfill(
        self,
        mailbox: str,
        *,
        newest_first: bool = True,
        batch_size: int = 200,
        since: date | None = None,
        before: date | None = None,
    ) -> AsyncIterator[list[RemoteMessage]]:
        """Yield message batches. ``since``/``before`` bound the import by IMAP
        internal date (``SINCE`` inclusive, ``BEFORE`` exclusive) so the caller
        can sync a recent window and/or the older remainder (SPEC §12 M2)."""
        info = await self._select(mailbox)
        self._uidvalidity_seen[mailbox] = info.get("uidvalidity", 0)
        imap = await self._connect()
        search = await imap.uid_search(_date_criteria(since, before))
        uids = parse_search_uids(search.lines)
        uids.sort(reverse=newest_first)
        for start in range(0, len(uids), batch_size):
            batch = uids[start : start + batch_size]
            messages = await self._fetch_messages(imap, batch)
            if messages:
                yield messages

    async def message_count(self, mailbox: str, *, since: date | None = None) -> int:
        """Number of messages to import for the progress denominator. Uses the
        cheap ``EXISTS`` from SELECT for the whole mailbox, or a dated SEARCH."""
        info = await self._select(mailbox)
        if since is None:
            return int(info.get("exists", 0))
        imap = await self._connect()
        search = await imap.uid_search(_date_criteria(since, None))
        return len(parse_search_uids(search.lines))

    async def newest_uid(self, mailbox: str) -> int:
        """Current high-water UID (``UIDNEXT - 1``); 0 for an empty mailbox. Used
        by the ``new`` sync window to skip history and only pull future mail."""
        info = await self._select(mailbox)
        return max(0, int(info.get("uidnext", 1)) - 1)

    async def list_uids(self, mailbox: str) -> set[int]:
        """Full set of UIDs currently present in the mailbox (expunge reconciliation).

        Cheap — a single ``UID SEARCH ALL``. The CONDSTORE poll only reports new and
        changed messages, not vanished ones (that needs QRESYNC/VANISHED); callers use
        this to detect server-side expunges for small folders like Drafts.
        """
        await self._select(mailbox)
        imap = await self._connect()
        resp = await imap.uid_search("ALL")
        return set(parse_search_uids(resp.lines))

    async def fetch_raw(self, mailbox: str, uid: int) -> bytes | None:
        """One message's raw bytes, or None if it is not (or no longer) there."""
        await self._select(mailbox)
        imap = await self._connect()
        for message in await self._fetch_messages(imap, [uid]):
            if message.uid == uid:
                return message.raw
        return None

    async def _fetch_messages(self, imap: aioimaplib.IMAP4, uids: list[int]) -> list[RemoteMessage]:
        if not uids:
            return []
        uid_set = ",".join(str(u) for u in uids)
        items = _FETCH_ITEMS_GM if self._gm_ext else _FETCH_ITEMS
        resp = await imap.uid("fetch", uid_set, items)
        if resp.result != "OK":
            raise ProviderConnectError("FETCH failed")
        out: list[RemoteMessage] = []
        for pf in parse_fetch_lines(resp.lines):
            if pf.raw is None or pf.internal_date is None:
                continue
            out.append(
                RemoteMessage(
                    uid=pf.uid,
                    raw=pf.raw,
                    flags=pf.flags,
                    internal_date=pf.internal_date,
                    gmail_labels=pf.gmail_labels,
                )
            )
        return out

    async def poll(
        self, mailbox: str, *, last_seen_uid: int, highestmodseq: int | None
    ) -> MailboxDelta:
        info = await self._select(mailbox)
        imap = await self._connect()
        current_uidvalidity = info.get("uidvalidity", 0)
        new_uidvalidity: int | None = None
        prev = self._uidvalidity_seen.get(mailbox)
        if prev is not None and current_uidvalidity != prev:
            new_uidvalidity = current_uidvalidity
        self._uidvalidity_seen[mailbox] = current_uidvalidity
        if new_uidvalidity is not None:
            return MailboxDelta([], [], [], new_uidvalidity, info.get("highestmodseq"))

        # New messages: UID SEARCH beyond last_seen_uid (filter the M:* boundary echo).
        search = await imap.uid_search(f"UID {last_seen_uid + 1}:*")
        candidate_uids = [u for u in parse_search_uids(search.lines) if u > last_seen_uid]
        new_messages = await self._fetch_messages(imap, sorted(candidate_uids))
        new_uid_set = {m.uid for m in new_messages}

        flag_changes: list[FlagChange] = []
        expunged_uids: list[int] = []
        if highestmodseq is not None and self._qresync_enabled:
            # QRESYNC fast resync (RFC 7162): a re-SELECT with (QRESYNC (uidvalidity
            # modseq)) makes the server report, in one pass, the FLAGS changed and the
            # UIDs VANISHED since ``modseq``. VANISHED is only surfaced on a *sync*
            # command, so we harvest both from the SELECT response here.
            qr = await self._qresync_select(mailbox, current_uidvalidity, highestmodseq)
            expunged_uids = [u for u in parse_vanished(qr.lines) if u not in new_uid_set]
            for pf in parse_fetch_lines(qr.lines):
                if pf.uid not in new_uid_set:
                    flag_changes.append(FlagChange(uid=pf.uid, flags=pf.flags))
        elif highestmodseq is not None:
            # CONDSTORE-only server (e.g. Gmail): flag delta via CHANGEDSINCE. There
            # is no VANISHED here — the engine reconciles expunges by full-UID diff.
            fresp = await imap.uid("fetch", "1:*", f"(UID FLAGS) (CHANGEDSINCE {highestmodseq})")
            if fresp.result == "OK":
                for pf in parse_fetch_lines(fresp.lines):
                    if pf.uid not in new_uid_set:
                        flag_changes.append(FlagChange(uid=pf.uid, flags=pf.flags))

        return MailboxDelta(
            new=new_messages,
            flag_changes=flag_changes,
            expunged_uids=expunged_uids,
            new_uidvalidity=None,
            new_highestmodseq=info.get("highestmodseq"),
        )

    async def _qresync_select(
        self, mailbox: str, uidvalidity: int, modseq: int
    ) -> aioimaplib.Response:
        """Raw ``SELECT <mailbox> (QRESYNC (uidvalidity modseq))`` (RFC 7162).

        Issued as a low-level command because aioimaplib's high-level ``select``
        takes no parameters, and because SELECT is a *sync* command — the only kind
        whose untagged responses (VANISHED, changed FETCH) aioimaplib buffers."""
        imap = await self._connect()
        proto = imap.protocol
        param = f"(QRESYNC ({uidvalidity} {modseq}))"
        cmd = aioimaplib.Command(
            "SELECT", proto.new_tag(), self._wire(mailbox), param, loop=proto.loop
        )
        return await proto.execute(cmd)

    async def idle(self, mailbox: str) -> None:
        imap = await self._connect()
        if "IDLE" not in (imap.protocol.capabilities or set()):
            raise ProviderNotSupported("server does not support IDLE")
        await self._select(mailbox)
        idle_task = await imap.idle_start(timeout=_IDLE_TIMEOUT)
        try:
            # No word from the server even after IDLE's own end: just return, and
            # the caller's poll shows whether the connection is still there.
            with contextlib.suppress(TimeoutError):
                await imap.wait_server_push(timeout=_IDLE_TIMEOUT + 60)
        finally:
            imap.idle_done()
            with contextlib.suppress(Exception):
                await asyncio.wait_for(idle_task, timeout=10)

    async def push_flags(
        self, mailbox: str, uid: int, add: frozenset[str], remove: frozenset[str]
    ) -> None:
        await self._select(mailbox)
        imap = await self._connect()
        if add:
            await imap.uid("store", str(uid), "+FLAGS.SILENT", f"({' '.join(sorted(add))})")
        if remove:
            await imap.uid("store", str(uid), "-FLAGS.SILENT", f"({' '.join(sorted(remove))})")

    async def create_mailbox(self, path: str) -> None:
        """Create a mailbox on the server (idempotent — ignores 'already exists')."""
        imap = await self._connect()
        with contextlib.suppress(Exception):
            await imap.create(self._wire(path))

    async def move(self, uid: int, src: str, dst: str) -> int:
        await self._select(src)
        imap = await self._connect()
        resp = await imap.uid("move", str(uid), self._wire(dst))
        if resp.result != "OK":
            raise ProviderConnectError(f"MOVE {uid} -> {dst} failed")
        for line in resp.lines:
            if isinstance(line, (bytes, bytearray)):
                m = _COPYUID_RE.search(bytes(line))
                if m:
                    return int(m.group(1))
        return uid

    async def delete_messages(self, mailbox: str, uids: list[int]) -> None:
        """Flag the given UIDs \\Deleted and EXPUNGE them from ``mailbox``."""
        if not uids:
            return
        await self._select(mailbox)
        imap = await self._connect()
        uid_set = ",".join(str(u) for u in uids)
        await imap.uid("store", uid_set, "+FLAGS.SILENT", "(\\Deleted)")
        try:
            # UID EXPUNGE (UIDPLUS) removes only our UIDs; fall back to a plain
            # EXPUNGE (all \Deleted in the mailbox) when unsupported.
            await imap.uid("expunge", uid_set)
        except Exception:
            await imap.expunge()

    async def append(self, mailbox: str, raw: bytes, flags: frozenset[str]) -> int:
        imap = await self._connect()
        flag_str = f"({' '.join(sorted(flags))})" if flags else None
        resp = await imap.append(raw, mailbox=self._wire(mailbox), flags=flag_str)
        if resp.result != "OK":
            raise ProviderConnectError(f"APPEND to {mailbox} failed")
        for line in resp.lines:
            if isinstance(line, (bytes, bytearray)):
                m = _APPENDUID_RE.search(bytes(line))
                if m:
                    return int(m.group(1))
        return 0

    async def send(self, raw: bytes, envelope_from: str, envelope_to: list[str]) -> None:
        cfg = self._cfg
        if cfg.public_only:
            await _require_public(cfg.smtp_host)
        use_tls = cfg.smtp_tls == "ssl"
        start_tls = cfg.smtp_tls == "starttls"
        host_slot = (
            contextlib.nullcontext()
            if cfg.internal
            else loop_semaphore("smtp_host", cfg.smtp_host, get_settings().per_host_conn_cap)
        )
        try:
            async with host_slot:  # bound simultaneous sockets to the SMTP host (§4)
                await aiosmtplib.send(
                    raw,
                    sender=envelope_from,
                    recipients=envelope_to,
                    hostname=cfg.smtp_host,
                    port=cfg.smtp_port,
                    username=cfg.username,
                    password=cfg.password,
                    use_tls=use_tls,
                    start_tls=start_tls,
                    tls_context=cfg.tls_context,
                    timeout=_CONNECT_TIMEOUT,
                    # Our own name, not the container's random host name.
                    local_hostname=smtp_helo_name(),
                )
        except aiosmtplib.SMTPException as exc:
            raise _smtp_error(exc) from exc