webmail-client/backend/tests/plugins/test_antispam_learning.py

"""C2 learning loop: what teaches, what doesn't, and what only teaches once.

Passive learning (SPEC §C.5) reads the user's own filing as a verdict. The value
is in the restraint — most folder moves say nothing about spam, and teaching the
same message twice skews the corpus.
"""

from __future__ import annotations

import uuid
from typing import Any

import pytest
from webmail_plugin_antispam import hook, jobs, policy

SETTINGS = policy.Settings(
    enabled_accounts=[],
    move_threshold=12.0,
    banner_threshold=6.0,
    network_checks=False,
    learn_from_provider_junk=True,
)


class _FakeSdk:
    def __init__(self, learned: str | None = None) -> None:
        self.learned_before = learned
        self.published: list[tuple[str, dict[str, Any]]] = []
        self.recorded: list[dict[str, Any]] = []
        self.enqueued: list[dict[str, Any]] = []

        outer = self

        class _M:
            async def envelope(self, *_a: object, **_k: object) -> dict[str, Any]:
                return {"account_id": "acc-1"}

            async def raw(self, *_a: object, **_k: object) -> bytes:
                return b"From: x@y.test\r\n\r\nbody\r\n"

        self.messages = _M()

        async def _enqueue(_f: object, **kw: Any) -> None:
            outer.enqueued.append(kw)

        self.enqueue = _enqueue

    async def publish(self, event: str, payload: dict[str, Any]) -> None:
        self.published.append((event, payload))


@pytest.fixture
def wired(monkeypatch: pytest.MonkeyPatch):  # type: ignore[no-untyped-def]
    def _install(*, learned: str | None = None, engine_up: bool = True) -> _FakeSdk:
        fake = _FakeSdk(learned)
        taught: list[bool] = []

        async def _learn(_raw: bytes, *, user_key: str, spam: bool) -> bool:
            taught.append(spam)
            return engine_up

        async def _already(_u: object, _m: object) -> str | None:
            return fake.learned_before

        async def _remember(_u: object, _m: object, **kw: Any) -> None:
            fake.recorded.append(kw)

        async def _settings(_u: object) -> policy.Settings:
            return SETTINGS

        monkeypatch.setattr(hook, "_SDK", fake)
        monkeypatch.setattr(hook.rspamd, "learn", _learn)
        monkeypatch.setattr(hook, "_already_learned", _already)
        monkeypatch.setattr(hook, "_remember_learned", _remember)
        monkeypatch.setattr(policy, "load_settings", _settings)
        fake.taught = taught  # type: ignore[attr-defined]
        return fake

    return _install


async def _move(before: str | None, now: str | None) -> None:
    await hook.on_message_changed(
        {
            "user_id": str(uuid.uuid4()),
            "message_id": str(uuid.uuid4()),
            "previous_special_use": before,
            "special_use": now,
        }
    )


async def test_filing_into_spam_teaches_spam(wired) -> None:  # type: ignore[no-untyped-def]
    fake = wired()
    await _move("inbox", "junk")
    assert fake.taught == [True]
    assert fake.recorded[0]["verdict"] == "spam"


async def test_rescuing_from_spam_teaches_ham(wired) -> None:  # type: ignore[no-untyped-def]
    fake = wired()
    await _move("junk", "inbox")
    assert fake.taught == [False]
    assert fake.recorded[0]["verdict"] == "ham"


@pytest.mark.parametrize(
    ("before", "now"),
    [
        ("inbox", "archive"),  # tidying up
        ("junk", "trash"),  # deleting spam is not a verdict about ham
        ("archive", "inbox"),  # unarchiving says nothing about spam
        ("sent", "archive"),
        ("inbox", "inbox"),  # not a move at all
    ],
)
async def test_other_moves_teach_nothing(wired, before, now) -> None:  # type: ignore[no-untyped-def]
    fake = wired()
    await _move(before, now)
    assert fake.taught == []
    assert fake.recorded == []


async def test_the_same_lesson_is_not_taught_twice(wired) -> None:  # type: ignore[no-untyped-def]
    """Users move things back and forth; the corpus should not count that twice."""
    fake = wired(learned="spam")
    await _move("inbox", "junk")
    assert fake.taught == []


async def test_reversing_a_lesson_does_teach(wired) -> None:  # type: ignore[no-untyped-def]
    """Having taught spam, moving it back out must be able to teach ham."""
    fake = wired(learned="spam")
    await _move("junk", "inbox")
    assert fake.taught == [False]


async def test_nothing_is_recorded_when_the_engine_is_down(wired) -> None:  # type: ignore[no-untyped-def]
    """`learned` stays unset so a later move can teach it again."""
    fake = wired(engine_up=False)
    await _move("inbox", "junk")
    assert fake.taught == [True]
    assert fake.recorded == []


# --- bootstrap ---------------------------------------------------------------


async def test_going_live_enqueues_the_bootstrap(wired) -> None:  # type: ignore[no-untyped-def]
    fake = wired()
    await hook.on_account_state_changed(
        {"user_id": str(uuid.uuid4()), "account_id": "acc-1", "state": "live"}
    )
    assert fake.enqueued == [{"user_id": fake.enqueued[0]["user_id"], "account_id": "acc-1"}]


async def test_other_states_do_not_enqueue(wired) -> None:  # type: ignore[no-untyped-def]
    fake = wired()
    for state in ("backfilling", "error", "paused"):
        await hook.on_account_state_changed(
            {"user_id": str(uuid.uuid4()), "account_id": "acc-1", "state": state}
        )
    assert fake.enqueued == []


async def test_bootstrap_respects_the_setting(monkeypatch: pytest.MonkeyPatch) -> None:
    """Opting out of provider-Junk learning means no training pass at all."""
    sampled: list[str] = []

    async def _settings(_u: object) -> policy.Settings:
        return policy.Settings(
            enabled_accounts=[],
            move_threshold=12.0,
            banner_threshold=6.0,
            network_checks=False,
            learn_from_provider_junk=False,
        )

    class _M:
        async def sample(self, *_a: object, **kw: object) -> list[uuid.UUID]:
            sampled.append(str(kw.get("special_use")))
            return []

    class _Fake:
        messages = _M()

    monkeypatch.setattr(policy, "load_settings", _settings)
    monkeypatch.setattr(jobs, "_SDK", _Fake())
    await jobs.bootstrap_account({}, user_id=str(uuid.uuid4()), account_id=str(uuid.uuid4()))
    assert sampled == []


async def test_bootstrap_refuses_a_lopsided_corpus(monkeypatch: pytest.MonkeyPatch) -> None:
    """Three spam beside a thousand ham would teach "everything is ham".

    Accounts do not sync the provider's spam folder by default, so this is the
    common case, not an edge case.
    """
    taught: list[bool] = []

    async def _settings(_u: object) -> policy.Settings:
        return SETTINGS

    class _Kv:
        async def get(self, _k: str, _d: object = None) -> list[str]:
            return []

        async def set(self, _k: str, _v: object) -> None:
            raise AssertionError("must not mark the account done")

    class _M:
        async def sample(self, *_a: object, **kw: object) -> list[uuid.UUID]:
            n = 3 if kw.get("special_use") == "junk" else 1000
            return [uuid.uuid4() for _ in range(n)]

        async def raw(self, *_a: object, **_k: object) -> bytes:
            return b"x"

    class _Fake:
        messages = _M()

        def kv(self, _u: object) -> _Kv:
            return _Kv()

    async def _learn(_raw: bytes, **kw: object) -> bool:
        taught.append(bool(kw.get("spam")))
        return True

    monkeypatch.setattr(policy, "load_settings", _settings)
    monkeypatch.setattr(jobs, "_SDK", _Fake())
    monkeypatch.setattr(jobs.rspamd, "learn", _learn)

    await jobs.bootstrap_account({}, user_id=str(uuid.uuid4()), account_id=str(uuid.uuid4()))
    assert taught == [], "taught from a corpus too lopsided to learn from"


async def test_bootstrap_caps_ham_at_a_multiple_of_spam(monkeypatch: pytest.MonkeyPatch) -> None:
    marked: list[object] = []

    async def _settings(_u: object) -> policy.Settings:
        return SETTINGS

    class _Kv:
        async def get(self, _k: str, _d: object = None) -> list[str]:
            return []

        async def set(self, _k: str, v: object) -> None:
            marked.append(v)

    class _M:
        async def sample(self, *_a: object, **kw: object) -> list[uuid.UUID]:
            n = 20 if kw.get("special_use") == "junk" else 1000
            return [uuid.uuid4() for _ in range(n)]

        async def raw(self, *_a: object, **_k: object) -> bytes:
            return b"x"

    class _Fake:
        messages = _M()
        published: list[object] = []

        def kv(self, _u: object) -> _Kv:
            return _Kv()

        async def publish(self, _e: str, _p: object) -> None:
            return None

    taught: list[bool] = []

    async def _learn(_raw: bytes, **kw: object) -> bool:
        taught.append(bool(kw.get("spam")))
        return True

    monkeypatch.setattr(policy, "load_settings", _settings)
    monkeypatch.setattr(jobs, "_SDK", _Fake())
    monkeypatch.setattr(jobs.rspamd, "learn", _learn)
    monkeypatch.setattr(jobs, "PAUSE_SECONDS", 0)

    await jobs.bootstrap_account({}, user_id=str(uuid.uuid4()), account_id=str(uuid.uuid4()))
    assert taught.count(True) == 20
    assert taught.count(False) == 60, "ham should be capped at 3x the spam count"
    assert marked, "a completed bootstrap must mark the account done"


# --- the sweep: bootstrapping must not depend on catching an event -----------


async def test_sweep_visits_every_live_account(monkeypatch: pytest.MonkeyPatch) -> None:
    """An account that went live before the plugin existed still gets a warm start."""
    visited: list[tuple[str, str]] = []
    accounts = [(uuid.uuid4(), uuid.uuid4()) for _ in range(3)]

    class _Accounts:
        async def live(self) -> list[tuple[uuid.UUID, uuid.UUID]]:
            return accounts

    class _Fake:
        accounts = _Accounts()

    async def _one(_ctx: object, *, user_id: str, account_id: str) -> None:
        visited.append((user_id, account_id))

    monkeypatch.setattr(jobs, "_SDK", _Fake())
    monkeypatch.setattr(jobs, "bootstrap_account", _one)
    await jobs.bootstrap_sweep({})
    assert visited == [(str(u), str(a)) for u, a in accounts]


async def test_one_failing_account_does_not_stop_the_sweep(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    done: list[str] = []
    accounts = [(uuid.uuid4(), uuid.uuid4()) for _ in range(3)]

    class _Accounts:
        async def live(self) -> list[tuple[uuid.UUID, uuid.UUID]]:
            return accounts

    class _Fake:
        accounts = _Accounts()

    async def _one(_ctx: object, *, user_id: str, account_id: str) -> None:
        if account_id == str(accounts[1][1]):
            raise RuntimeError("that mailbox is on fire")
        done.append(account_id)

    monkeypatch.setattr(jobs, "_SDK", _Fake())
    monkeypatch.setattr(jobs, "bootstrap_account", _one)
    await jobs.bootstrap_sweep({})
    assert done == [str(accounts[0][1]), str(accounts[2][1])]


@pytest.mark.asyncio
async def test_a_dead_engine_never_marks_the_account_bootstrapped(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """An engine that refuses the FIRST message must not look like a clean run.

    This shipped: failure was signalled by negating the count taught so far, so
    a refusal on message one returned `-0` — which is `0`, not negative. The
    caller read that as "learned nothing, fine", marked the account done, and
    the hourly sweep then skipped it forever. The account most in need of a warm
    start was the one guaranteed never to get one.
    """
    marked: list[object] = []

    async def _settings(_u: object) -> policy.Settings:
        return SETTINGS

    class _Kv:
        async def get(self, _k: str, _d: object = None) -> list[str]:
            return []

        async def set(self, _k: str, v: object) -> None:
            marked.append(v)

    class _M:
        async def sample(self, *_a: object, **kw: object) -> list[uuid.UUID]:
            n = 50 if kw.get("special_use") == "junk" else 50
            return [uuid.uuid4() for _ in range(n)]

        async def raw(self, *_a: object, **_k: object) -> bytes:
            return b"x"

    class _Fake:
        messages = _M()

        def kv(self, _u: object) -> _Kv:
            return _Kv()

        async def publish(self, *_a: object, **_k: object) -> None:
            """Present so this test fails on its assertion, not on a bare fake."""

    async def _dead(_raw: bytes, **_kw: object) -> bool:
        return False  # every learn refused, from the very first one

    monkeypatch.setattr(policy, "load_settings", _settings)
    monkeypatch.setattr(jobs, "_SDK", _Fake())
    monkeypatch.setattr(jobs.rspamd, "learn", _dead)

    await jobs.bootstrap_account({}, user_id=str(uuid.uuid4()), account_id=str(uuid.uuid4()))
    assert marked == [], "marked the account bootstrapped after learning nothing"


@pytest.mark.asyncio
async def test_an_engine_that_dies_midway_also_leaves_the_account_unmarked(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """The case the original sentinel did handle — keep it handled."""
    marked: list[object] = []
    calls = {"n": 0}

    async def _settings(_u: object) -> policy.Settings:
        return SETTINGS

    class _Kv:
        async def get(self, _k: str, _d: object = None) -> list[str]:
            return []

        async def set(self, _k: str, v: object) -> None:
            marked.append(v)

    class _M:
        async def sample(self, *_a: object, **_kw: object) -> list[uuid.UUID]:
            return [uuid.uuid4() for _ in range(50)]

        async def raw(self, *_a: object, **_k: object) -> bytes:
            return b"x"

    class _Fake:
        messages = _M()

        def kv(self, _u: object) -> _Kv:
            return _Kv()

        async def publish(self, *_a: object, **_k: object) -> None:
            """As above."""

    async def _dies(_raw: bytes, **_kw: object) -> bool:
        calls["n"] += 1
        return calls["n"] <= 5

    monkeypatch.setattr(policy, "load_settings", _settings)
    monkeypatch.setattr(jobs, "_SDK", _Fake())
    monkeypatch.setattr(jobs.rspamd, "learn", _dies)

    await jobs.bootstrap_account({}, user_id=str(uuid.uuid4()), account_id=str(uuid.uuid4()))
    assert marked == [], "marked the account bootstrapped after a partial run"


def test_compose_gives_the_app_the_controller_password() -> None:
    """The learning loop needs the controller password; classification does not.

    Rspamd splits over two ports: the scan worker takes /checkv2 anonymously,
    the controller demands a password for /stat and /learnspam. Ship the password
    to rspamd but not to the api/worker containers and the result is a filter
    that looks healthy — mail is still classified — while every learn is refused
    401 and swallowed by the plugin's fail-open path. The only visible trace is
    `"bayes": null` on the stats endpoint.

    That shipped, so it is worth a test that reads the compose files rather than
    trusting review to catch the missing line again.
    """
    import pathlib

    root = pathlib.Path(__file__).resolve().parents[3]
    for rel, anchor in (
        ("docker-compose.yml", "x-backend-env:"),
        ("docker/docker-compose.yml", "environment: &worker-env"),
    ):
        text = (root / rel).read_text()
        assert anchor in text, f"{rel}: shared env block moved; update this test"
        start = text.index(anchor)
        # The shared block runs to the next top-level key.
        tail = text[start:]
        end = start + (tail.index("\nservices:") if "\nservices:" in tail else len(tail))
        block = text[start:end]
        assert "RSPAMD_CONTROLLER_PASSWORD" in block, (
            f"{rel}: the app containers do not receive RSPAMD_CONTROLLER_PASSWORD — "
            "antispam will classify but never learn"
        )