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

"""Pure parsing helpers for IMAP responses (unit-testable, no I/O).

aioimaplib returns responses as lists of ``bytes``/``bytearray`` lines. A FETCH
with ``BODY[]`` yields, per message, a header line ending ``... {<size>}`` followed
by a ``bytearray`` literal, then a lone ``)``. LIST lines carry SPECIAL-USE flags.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from datetime import datetime

from imapclient.imap_utf7 import decode as imap_utf7_decode

# --- SPECIAL-USE mapping (SPEC §4 special_use values) -----------------------

_SPECIAL_USE_FLAGS = {
    "\\sent": "sent",
    "\\drafts": "drafts",
    "\\trash": "trash",
    "\\junk": "junk",
    "\\archive": "archive",
    "\\all": "archive",  # Gmail "All Mail" acts as archive
}

_LIST_RE = re.compile(rb'^\(([^)]*)\)\s+(?:"([^"]*)"|NIL)\s+(.+)$')
_UID_RE = re.compile(rb"UID\s+(\d+)")
_FLAGS_RE = re.compile(rb"FLAGS\s+\(([^)]*)\)")
_INTERNALDATE_RE = re.compile(rb'INTERNALDATE\s+"([^"]+)"')
_LITERAL_RE = re.compile(rb"\{(\d+)\}\s*$")
_FETCH_HEADER_RE = re.compile(rb"^\d+\s+FETCH\s+\(", re.IGNORECASE)


@dataclass(frozen=True)
class ParsedMailbox:
    path: str
    special_use: str | None


@dataclass(frozen=True)
class ParsedFetch:
    uid: int
    flags: frozenset[str]
    internal_date: datetime | None
    raw: bytes | None  # None when the fetch did not request the body


def _decode_mailbox_name(raw: bytes) -> str:
    name = raw.strip()
    if name.startswith(b'"') and name.endswith(b'"'):
        name = name[1:-1]
    # IMAP mailbox names are modified UTF-7 (SPEC §4: store UTF-7 decoded).
    return str(imap_utf7_decode(bytes(name)))


def parse_list_line(line: bytes) -> ParsedMailbox | None:
    """Parse one LIST response line into a :class:`ParsedMailbox` (or None)."""
    m = _LIST_RE.match(line.strip())
    if not m:
        return None
    flags = m.group(1).lower().split()
    name = _decode_mailbox_name(m.group(3))
    special: str | None = None
    for flag in flags:
        mapped = _SPECIAL_USE_FLAGS.get(flag.decode("ascii", "ignore"))
        if mapped:
            special = mapped
            break
    if special is None and name.upper() == "INBOX":
        special = "inbox"
    return ParsedMailbox(path=name, special_use=special)


def parse_internaldate(raw: bytes | str) -> datetime | None:
    """Parse an IMAP INTERNALDATE (``d-Mon-Y H:M:S +ZZZZ``) to an aware datetime."""
    text = raw.decode("ascii", "ignore") if isinstance(raw, bytes) else raw
    text = text.strip()
    try:
        return datetime.strptime(text, "%d-%b-%Y %H:%M:%S %z")
    except ValueError:
        return None


def _parse_flag_token(tok: bytes) -> str:
    return tok.decode("ascii", "ignore")


def parse_fetch_header(line: bytes) -> tuple[int | None, frozenset[str], datetime | None]:
    """Extract (uid, flags, internaldate) from a FETCH header line."""
    uid_m = _UID_RE.search(line)
    uid = int(uid_m.group(1)) if uid_m else None
    flags_m = _FLAGS_RE.search(line)
    if flags_m:
        flags = frozenset(_parse_flag_token(t) for t in flags_m.group(1).split())
    else:
        flags = frozenset()
    date_m = _INTERNALDATE_RE.search(line)
    internal = parse_internaldate(date_m.group(1)) if date_m else None
    return uid, flags, internal


def parse_fetch_lines(lines: list[bytes | bytearray]) -> list[ParsedFetch]:
    """Walk a full FETCH response, returning one :class:`ParsedFetch` per message.

    Handles both metadata-only fetches (FLAGS/INTERNALDATE) and body fetches
    (a ``{size}`` literal announced on the header line, body on the next line).
    """
    results: list[ParsedFetch] = []
    i = 0
    n = len(lines)
    while i < n:
        line = lines[i]
        if not isinstance(line, (bytes, bytearray)) or not _FETCH_HEADER_RE.match(bytes(line)):
            i += 1
            continue
        header = bytes(line)
        uid, flags, internal = parse_fetch_header(header)
        raw: bytes | None = None
        if _LITERAL_RE.search(header) and i + 1 < n:
            body = lines[i + 1]
            raw = bytes(body)
            i += 2
            # Skip the trailing ")" continuation line if present.
            if (
                i < n
                and isinstance(lines[i], (bytes, bytearray))
                and bytes(lines[i]).strip() == b")"
            ):
                i += 1
        else:
            i += 1
        if uid is not None:
            results.append(ParsedFetch(uid=uid, flags=flags, internal_date=internal, raw=raw))
    return results


def parse_search_uids(lines: list[bytes | bytearray]) -> list[int]:
    """Parse a UID SEARCH response's first line into a list of ints."""
    if not lines:
        return []
    first = lines[0]
    if not isinstance(first, (bytes, bytearray)):
        return []
    return [int(tok) for tok in bytes(first).split() if tok.isdigit()]


def parse_select_untagged(lines: list[bytes | bytearray]) -> dict[str, int]:
    """Extract UIDVALIDITY / UIDNEXT / HIGHESTMODSEQ / EXISTS from SELECT lines."""
    out: dict[str, int] = {}
    patterns = {
        "uidvalidity": re.compile(rb"UIDVALIDITY\s+(\d+)"),
        "uidnext": re.compile(rb"UIDNEXT\s+(\d+)"),
        "highestmodseq": re.compile(rb"HIGHESTMODSEQ\s+(\d+)"),
        "exists": re.compile(rb"^(\d+)\s+EXISTS"),
    }
    for line in lines:
        if not isinstance(line, (bytes, bytearray)):
            continue
        b = bytes(line)
        for key, pat in patterns.items():
            m = pat.search(b)
            if m:
                out[key] = int(m.group(1))
    return out