webmail-client/plugins/webmail_plugin_antispam/webmail_plugin_antispam/hook.py

"""``message.new`` hook — the C1 classification flow (SPEC §C.4).

Runs in the sync worker wherever ``message.new`` is published. The invariant that
shapes every branch: **spam filtering must never delay or lose mail**. Every
failure path here ends in "deliver normally, unclassified" — an unreachable
engine, a missing blob, an account with no Junk folder, a database error while
recording the verdict.
"""

from __future__ import annotations

import json
import logging
import uuid
from typing import Any

from sqlalchemy import text

from app.plugins.sdk import sdk
from webmail_plugin_antispam import policy, rspamd

logger = logging.getLogger("plg.antispam")

_SDK = sdk("antispam")


async def _record(
    user_id: uuid.UUID,
    message_id: uuid.UUID,
    *,
    score: float,
    symbols: list[str],
    action: str,
) -> None:
    """Write the verdict row. A failure here must not affect delivery."""
    try:
        async with _SDK.session(user_id) as s:
            await s.execute(
                text(
                    "INSERT INTO plg_antispam_verdicts "
                    "  (message_id, user_id, score, symbols, action) "
                    "VALUES (:m, :u, :sc, CAST(:sy AS jsonb), :a) "
                    "ON CONFLICT (message_id) DO NOTHING"
                ),
                {
                    "m": str(message_id),
                    "u": str(user_id),
                    "sc": score,
                    "sy": json.dumps(symbols),
                    "a": action,
                },
            )
            await s.commit()
    except Exception as exc:  # noqa: BLE001
        logger.warning("could not record antispam verdict for %s: %s", message_id, exc)


async def on_message_new(payload: dict[str, Any]) -> None:
    uid_raw = payload.get("user_id")
    mid_raw = payload.get("message_id")
    if not isinstance(uid_raw, str) or not isinstance(mid_raw, str):
        return
    user_id, message_id = uuid.UUID(uid_raw), uuid.UUID(mid_raw)

    env = await _SDK.messages.envelope(user_id, message_id)
    if env is None:
        return

    # Only newly-arrived inbox mail (§C.4). Never Sent/Drafts/Junk/Trash: scoring
    # a message the user just sent, or one the provider already filed as spam,
    # would be noise at best and a misfiled reply at worst.
    if env.get("mailbox_special_use") != "inbox":
        return

    # Never during backfill. message.new fires for every imported message too, and
    # classifying history is wrong twice over: it would file years-old inbox mail
    # into Junk on the server, and it would do so inline in the backfill loop —
    # ~70 ms of HTTP per message, serialized, on an account with thousands of
    # them. It is also when accuracy is lowest, because Bayes has learned nothing
    # yet. The backfill window belongs to the C2 learning bootstrap (§C.5), not
    # to classification.
    if env.get("account_state") != "live":
        return

    settings = await policy.load_settings(user_id)
    account_id = str(env.get("account_id") or "")
    if not settings.covers(account_id):
        return

    rules = await policy.load_rules(user_id)
    sender = str(env.get("from_email") or "")

    # A deny-list hit needs no engine at all (acceptance 5), so check rules first
    # and only pay for scoring when the lists have nothing to say.
    action = policy.decide(sender=sender, rules=rules, score=None, settings=settings)
    score, symbols = 0.0, []

    if action == "none":  # neither list matched — ask the engine
        raw = await _SDK.messages.raw(user_id, message_id)
        if raw is not None:
            verdict = await rspamd.check(
                raw, user_key=str(user_id), network_checks=settings.network_checks
            )
            if verdict is not None:
                score, symbols = verdict.score, verdict.symbols
                action = policy.decide(
                    sender=sender, rules=rules, score=verdict.score, settings=settings
                )
                # Only ask the address book when the answer could change the
                # outcome — i.e. when we are about to move the message. On the
                # overwhelmingly common path (not spam) this costs nothing.
                if action == "moved" and await policy.sender_in_contacts(
                    user_id, sender
                ):
                    action = policy.decide(
                        sender=sender,
                        rules=rules,
                        score=verdict.score,
                        settings=settings,
                        contact=True,
                    )

    if action in ("moved", "deny_listed"):
        moved = await _SDK.actions.move_to_special(user_id, message_id, "junk")
        if not moved:
            # No Junk folder on this account (or it is already there). Downgrade
            # to a banner rather than dropping the verdict: the user still sees
            # the warning, and nothing is lost.
            action = "banner" if action == "moved" else "deny_listed"

    await _record(user_id, message_id, score=score, symbols=symbols, action=action)

    await _SDK.publish(
        "plugin.antispam.classified",
        {
            "user_id": str(user_id),
            "message_id": str(message_id),
            "score": score,
            "action": action,
        },
    )


async def _already_learned(user_id: uuid.UUID, message_id: uuid.UUID) -> str | None:
    async with _SDK.session(user_id) as s:
        row = (
            await s.execute(
                text(
                    "SELECT learned FROM plg_antispam_verdicts "
                    "WHERE message_id = :m AND user_id = :u"
                ),
                {"m": str(message_id), "u": str(user_id)},
            )
        ).first()
    return None if row is None else row[0]


async def _remember_learned(
    user_id: uuid.UUID, message_id: uuid.UUID, *, verdict: str, action: str
) -> None:
    """Record the teaching. Upsert, because the message may never have been
    classified by us — provider-filed spam the user drags back has no verdict row."""
    async with _SDK.session(user_id) as s:
        await s.execute(
            text(
                "INSERT INTO plg_antispam_verdicts "
                "  (message_id, user_id, score, symbols, action, learned) "
                "VALUES (:m, :u, 0, '[]'::jsonb, :a, :l) "
                "ON CONFLICT (message_id) DO UPDATE SET learned = EXCLUDED.learned"
            ),
            {"m": str(message_id), "u": str(user_id), "a": action, "l": verdict},
        )
        await s.commit()


async def on_message_changed(payload: dict[str, Any]) -> None:
    """Passive learning from where the user files things (SPEC §C.5).

    Moving mail into Spam says "this is spam" more honestly than any button,
    because the user is doing it for their own reasons rather than to train a
    filter. Moving it back out of Spam says the opposite.

    Only those two transitions teach. Spam -> Trash is a user tidying up, not a
    verdict; Inbox -> Archive says nothing about spam at all.
    """
    uid_raw, mid_raw = payload.get("user_id"), payload.get("message_id")
    if not isinstance(uid_raw, str) or not isinstance(mid_raw, str):
        return
    now, before = payload.get("special_use"), payload.get("previous_special_use")
    if now == before:
        return

    if now == "junk":
        verdict, action = "spam", "moved"
    elif before == "junk" and now == "inbox":
        verdict, action = "ham", "none"
    else:
        return

    user_id, message_id = uuid.UUID(uid_raw), uuid.UUID(mid_raw)

    # Teaching the same message the same thing twice skews the corpus, and users
    # do move things back and forth.
    if await _already_learned(user_id, message_id) == verdict:
        return

    settings = await policy.load_settings(user_id)
    env = await _SDK.messages.envelope(user_id, message_id)
    if env is None or not settings.covers(str(env.get("account_id") or "")):
        return

    raw = await _SDK.messages.raw(user_id, message_id)
    if raw is None:
        return
    if not await rspamd.learn(raw, user_key=str(user_id), spam=(verdict == "spam")):
        # Engine down: leave `learned` unset so a later move can teach it again.
        return

    await _remember_learned(user_id, message_id, verdict=verdict, action=action)
    await _SDK.publish(
        "plugin.antispam.learned",
        {"user_id": str(user_id), "message_id": str(message_id), "verdict": verdict},
    )


async def on_account_state_changed(payload: dict[str, Any]) -> None:
    """When an account finishes its initial import, warm Bayes from what it holds.

    This is the constructive use of the backfill that classification deliberately
    skips (§C.4): the same history we refuse to act on is exactly what we want to
    learn from.
    """
    if payload.get("state") != "live":
        return
    user_id, account_id = payload.get("user_id"), payload.get("account_id")
    if not isinstance(user_id, str) or not isinstance(account_id, str):
        return
    from webmail_plugin_antispam.jobs import bootstrap_account

    await _SDK.enqueue(bootstrap_account, user_id=user_id, account_id=account_id)