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

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 (
    parse_fetch_lines,
    parse_list_line,
    parse_search_uids,
    parse_select_untagged,
)

_FETCH_ITEMS = "(UID FLAGS INTERNALDATE RFC822.SIZE 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+)")


@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] = []

    # --- 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
        return imap

    async def _select(self, mailbox: str) -> dict[str, int]:
        imap = await self._connect()
        resp = await imap.select(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(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
    ) -> AsyncIterator[list[RemoteMessage]]:
        info = await self._select(mailbox)
        self._uidvalidity_seen[mailbox] = info.get("uidvalidity", 0)
        imap = await self._connect()
        search = await imap.uid_search("ALL")
        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 _fetch_messages(self, imap: aioimaplib.IMAP4, uids: list[int]) -> list[RemoteMessage]:
        if not uids:
            return []
        uid_set = ",".join(str(u) for u in uids)
        resp = await imap.uid("fetch", uid_set, _FETCH_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
                )
            )
        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))

        # Flag changes via CONDSTORE when we have a modseq baseline.
        flag_changes: list[FlagChange] = []
        if highestmodseq is not None:
            fresp = await imap.uid("fetch", "1:*", f"(UID FLAGS) (CHANGEDSINCE {highestmodseq})")
            if fresp.result == "OK":
                new_uid_set = {m.uid for m in new_messages}
                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))

        # TODO(spec): expunge detection needs QRESYNC/VANISHED (M2 hardening).
        return MailboxDelta(
            new=new_messages,
            flag_changes=flag_changes,
            expunged_uids=[],
            new_uidvalidity=None,
            new_highestmodseq=info.get("highestmodseq"),
        )

    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=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