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
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
from app.providers.base import (
    FlagChange,
    MailboxDelta,
    MailProvider,
    ProviderAuthError,
    ProviderConnectError,
    ProviderNotSupported,
    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[])"
_IDLE_TIMEOUT = 25 * 60  # SPEC §5.1: re-issue IDLE at least every 25 min
_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")


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"


@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


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
        )
        host = loop_semaphore("imap_host", cfg.imap_host, get_settings().per_host_conn_cap)
        await account.acquire()  # consistent order (account, then host) avoids deadlock
        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
        await self._acquire_conn_slots()
        try:
            try:
                if cfg.imap_tls == "ssl":
                    imap = aioimaplib.IMAP4_SSL(
                        host=cfg.imap_host, port=cfg.imap_port, timeout=_CONNECT_TIMEOUT
                    )
                else:
                    imap = aioimaplib.IMAP4(
                        host=cfg.imap_host, port=cfg.imap_port, timeout=_CONNECT_TIMEOUT
                    )
                await imap.wait_hello_from_server()
                if cfg.imap_tls == "starttls":
                    await imap.starttls()
            except (OSError, TimeoutError) as exc:
                raise ProviderConnectError(f"cannot connect to {cfg.imap_host}: {exc}") 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_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:
            await imap.wait_server_push()
        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(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), 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
        use_tls = cfg.smtp_tls == "ssl"
        start_tls = cfg.smtp_tls == "starttls"
        host_slot = 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,
                    timeout=_CONNECT_TIMEOUT,
                )
        except aiosmtplib.SMTPAuthenticationError as exc:
            raise ProviderAuthError(f"SMTP auth failed: {exc}") from exc
        except aiosmtplib.SMTPException as exc:
            raise ProviderConnectError(f"SMTP send failed: {exc}") from exc