191 lines
6.2 KiB
Python
191 lines
6.2 KiB
Python
"""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 == []
|