learningcircuit--local-deep-research/tests/security/test_egress_run_classification.py

"""Tests for the stage-B egress resolver/assembly/audit (ADR-0007).

Covers label resolution from declared attributes + dynamic refinements
(collection is_public, store URL fail-up), run assembly, and the audit-mode
decision (which must never raise).
"""

from __future__ import annotations

from unittest.mock import patch

from local_deep_research.security.egress import run_classification as rc
from local_deep_research.security.egress.classification import (
    Exposure,
    Label,
    Role,
    Sensitivity,
)
from local_deep_research.security.egress.policy import (
    EgressContext,
    EgressScope,
)

S = Sensitivity.SENSITIVE
NS = Sensitivity.NON_SENSITIVE
EXP = Exposure.EXPOSING
CON = Exposure.CONTAINED

POLICY = "local_deep_research.security.egress.policy"


def ctx(primary: str = "arxiv", username=None) -> EgressContext:
    return EgressContext(
        scope=EgressScope.BOTH,
        primary_engine=primary,
        require_local_llm=False,
        require_local_embeddings=False,
        username=username,
    )


class _PublicEngine:
    egress_sensitivity = NS
    egress_exposure = EXP


class _LocalStore:
    egress_sensitivity = S
    egress_exposure = CON
    url_setting = "search.engine.web.paperless.default_params.api_url"


# ---------------------------------------------------------------------------
# engine_label — static engines + URL fail-up
# ---------------------------------------------------------------------------


def test_public_engine_label():
    with patch(f"{POLICY}._get_engine_class", return_value=_PublicEngine):
        assert rc.engine_label("arxiv", {}, ctx()) == Label(NS, EXP)


def test_unknown_engine_label_is_none():
    with patch(f"{POLICY}._get_engine_class", return_value=None):
        assert rc.engine_label("mystery", {}, ctx()) is None


def test_local_store_contained_on_local_url():
    with (
        patch(f"{POLICY}._get_engine_class", return_value=_LocalStore),
        patch(f"{POLICY}._classify_engine_url", return_value=True),
    ):
        assert rc.engine_label("paperless", {}, ctx()) == Label(S, CON)


def test_local_store_fails_up_to_exposing_on_public_url():
    # Quadrant 4: sensitive data + exposing sink (public host).
    with (
        patch(f"{POLICY}._get_engine_class", return_value=_LocalStore),
        patch(f"{POLICY}._classify_engine_url", return_value=False),
    ):
        assert rc.engine_label("paperless", {}, ctx()) == Label(S, EXP)


# ---------------------------------------------------------------------------
# engine_label — collections / library (sensitivity from is_public)
# ---------------------------------------------------------------------------


def test_private_collection_is_sensitive():
    with patch(f"{POLICY}._resolve_collection_is_public", return_value=False):
        assert rc.engine_label("collection_abc", {}, ctx()) == Label(S, CON)


def test_public_collection_is_non_sensitive():
    with patch(f"{POLICY}._resolve_collection_is_public", return_value=True):
        assert rc.engine_label("collection_abc", {}, ctx()) == Label(NS, CON)


def test_library_aggregate_is_sensitive():
    with patch(f"{POLICY}._resolve_collection_is_public", return_value=False):
        assert rc.engine_label("library", {}, ctx()) == Label(S, CON)


# ---------------------------------------------------------------------------
# provider labels (inference sinks — exposure only)
# ---------------------------------------------------------------------------


def test_llm_labels():
    assert rc.llm_label("ollama") == Label(NS, CON)
    assert rc.llm_label("anthropic") == Label(NS, EXP)
    assert rc.llm_label("totally-unknown") == Label(NS, EXP)  # fail closed


def test_embeddings_labels():
    assert rc.embeddings_label("ollama") == Label(NS, CON)
    assert rc.embeddings_label("sentence_transformers") == Label(NS, CON)
    assert rc.embeddings_label("openai") == Label(NS, EXP)


# ---------------------------------------------------------------------------
# classify_run — assembly into SOURCE + SEARCH_SINK + INFERENCE_SINK
# ---------------------------------------------------------------------------


def test_classify_run_assembles_roles():
    with patch.object(rc, "engine_label", return_value=Label(NS, EXP)):
        comps = rc.classify_run(
            {}, ctx(), engines=["arxiv"], llm_provider="anthropic"
        )
    got = {(c.name, c.role) for c in comps}
    assert got == {
        ("arxiv", Role.SOURCE),
        ("arxiv", Role.SEARCH_SINK),
        ("llm:anthropic", Role.INFERENCE_SINK),
    }


def test_classify_run_fails_closed_on_unknown_engine():
    with patch.object(rc, "engine_label", return_value=None):
        comps = rc.classify_run(
            {}, ctx(), engines=["mystery"], llm_provider="ollama"
        )
    # Unknown engine is NOT dropped — it fails closed to sensitive+exposing.
    by = {(c.name, c.role): c for c in comps}
    assert by[("mystery", Role.SOURCE)].label == Label(S, EXP)
    assert ("mystery", Role.SEARCH_SINK) in by
    assert ("llm:ollama", Role.INFERENCE_SINK) in by


# ---------------------------------------------------------------------------
# audit_run — decision + never-raises contract
# ---------------------------------------------------------------------------


def test_audit_run_denies_sensitive_plus_cloud_llm():
    with patch.object(rc, "engine_label", return_value=Label(S, CON)):
        d = rc.audit_run(
            {}, ctx(), engines=["collection_x"], llm_provider="anthropic"
        )
    assert d is not None and not d.allowed
    assert d.reason == "sensitive_to_exposing_inference"


def test_audit_run_allows_public_plus_local_llm():
    with patch.object(rc, "engine_label", return_value=Label(NS, EXP)):
        d = rc.audit_run({}, ctx(), engines=["arxiv"], llm_provider="ollama")
    assert d is not None and d.allowed


def test_audit_run_permissive_under_unprotected_scope():
    # A normally-denied combo is ALLOWED (permissive) under the escape hatch,
    # with the diagnostic preserved.
    unprot = EgressContext(
        scope=EgressScope.UNPROTECTED,
        primary_engine="arxiv",
        require_local_llm=False,
        require_local_embeddings=False,
    )
    with patch.object(rc, "engine_label", return_value=Label(S, CON)):
        d = rc.audit_run(
            {}, unprot, engines=["collection_x"], llm_provider="anthropic"
        )
    assert d is not None and d.allowed
    assert d.reason.startswith("permissive:")


def test_audit_run_fails_closed_on_internal_error():
    # An uncomputable decision must refuse the run (fail closed) and never
    # raise — silent degradation to the scope PEPs would hide the failure.
    with patch.object(rc, "engine_label", side_effect=RuntimeError("boom")):
        d = rc.audit_run({}, ctx(), engines=["x"], llm_provider="ollama")
    assert d is not None and not d.allowed and d.reason == "audit_error"


def test_fail_closed_decision_needs_no_policy_import():
    # _fail_closed_decision runs from the except blocks that handle failures
    # (including a .policy import failure), so it must NOT import .policy. It
    # denies for a normal ctx and allows only under unprotected, using the
    # ctx's str-enum scope value directly.
    assert not rc._fail_closed_decision(ctx()).allowed
    unprot = EgressContext(
        scope=EgressScope.UNPROTECTED,
        primary_engine="arxiv",
        require_local_llm=False,
        require_local_embeddings=False,
    )
    assert rc._fail_closed_decision(unprot).allowed


def test_audit_run_internal_error_never_blocks_under_unprotected():
    # Fail-closed must not override the escape hatch: an internal error under
    # UNPROTECTED is still allowed.
    unprot = EgressContext(
        scope=EgressScope.UNPROTECTED,
        primary_engine="arxiv",
        require_local_llm=False,
        require_local_embeddings=False,
    )
    with patch.object(rc, "engine_label", side_effect=RuntimeError("boom")):
        d = rc.audit_run({}, unprot, engines=["x"], llm_provider="ollama")
    assert d is not None and d.allowed


def test_audit_run_unknown_engine_fails_closed_and_denies():
    # Unknown engine -> sensitive+exposing source; a cloud LLM -> exposing
    # inference sink -> the run is denied (not silently allowed).
    with patch.object(rc, "engine_label", return_value=None):
        d = rc.audit_run(
            {}, ctx(), engines=["mystery"], llm_provider="anthropic"
        )
    assert d is not None and not d.allowed
    assert d.reason == "sensitive_to_exposing_inference"


# ---------------------------------------------------------------------------
# audit_run_from_snapshot — shared route/worker entry point (covers all run
# entry points: API precheck, follow-up, chat, queue)
# ---------------------------------------------------------------------------


def test_audit_run_from_snapshot_denies_private_collection_plus_cloud_llm():
    # The worker-chokepoint case that a follow-up / chat / queue run reaches:
    # a private collection primary + a cloud LLM under a `both` scope must be
    # refused, matching the /api/start_research route.
    snap = {"llm.provider": "anthropic"}
    with patch.object(rc, "engine_label", return_value=Label(S, CON)):
        d = rc.audit_run_from_snapshot(
            snap, ctx(primary="collection_x"), "collection_x"
        )
    assert d is not None and not d.allowed
    assert d.reason == "sensitive_to_exposing_inference"


def test_audit_run_from_snapshot_unwraps_dict_wrapped_snapshot():
    # get_all_settings snapshots store each value as {"value": ...}; the helper
    # must unwrap it, else the LLM provider is dropped and the check is blind.
    snap = {"llm.provider": {"value": "anthropic"}}
    with patch.object(rc, "engine_label", return_value=Label(S, CON)):
        d = rc.audit_run_from_snapshot(
            snap, ctx(primary="collection_x"), "collection_x"
        )
    assert d is not None and not d.allowed


def test_audit_run_from_snapshot_permissive_under_unprotected():
    unprot = EgressContext(
        scope=EgressScope.UNPROTECTED,
        primary_engine="collection_x",
        require_local_llm=False,
        require_local_embeddings=False,
    )
    with patch.object(rc, "engine_label", return_value=Label(S, CON)):
        d = rc.audit_run_from_snapshot(
            {"llm.provider": "anthropic"}, unprot, "collection_x"
        )
    assert d is not None and d.allowed


def test_audit_run_from_snapshot_missing_llm_provider_uses_run_default():
    # A snapshot without llm.provider must resolve to the run's own default
    # (ollama, local) via get_setting_from_snapshot — the same value the run
    # would use — not be silently read as "no LLM". Sensitive source + local
    # default inference is quadrant 2 (allowed), matching the run's behaviour.
    with patch.object(rc, "engine_label", return_value=Label(S, CON)):
        d = rc.audit_run_from_snapshot(
            {}, ctx(primary="collection_x"), "collection_x"
        )
    assert d is not None and d.allowed


def test_audit_run_from_snapshot_includes_embeddings_for_collection():
    # A RAG run over a collection embeds, so a cloud embedder is an exposing
    # sink for the sensitive source -> refused even with a local LLM. The RAG
    # engine reads local_search_embedding_provider, so the audit must too.
    snap = {
        "llm.provider": "ollama",
        "local_search_embedding_provider": "openai",
    }
    with patch.object(rc, "engine_label", return_value=Label(S, CON)):
        d = rc.audit_run_from_snapshot(
            snap, ctx(primary="collection_x"), "collection_x"
        )
    assert d is not None and not d.allowed


def test_audit_run_from_snapshot_ignores_embeddings_for_lexical_store():
    # A non-collection sensitive store (paperless/elasticsearch) never embeds,
    # so a configured cloud embedder must NOT be pulled in: sensitive source +
    # local LLM is admissible (quadrant 2), not refused.
    snap = {
        "llm.provider": "ollama",
        "local_search_embedding_provider": "openai",
    }
    with patch.object(rc, "engine_label", return_value=Label(S, CON)):
        d = rc.audit_run_from_snapshot(
            snap, ctx(primary="paperless"), "paperless"
        )
    assert d is not None and d.allowed


def test_audit_run_from_snapshot_empty_llm_provider_resolves_from_snapshot():
    # Chat passes model_provider="" (not None) when no override is set. An empty
    # string must be treated as "resolve from snapshot", not passed through —
    # else classify_run's truthy guard silently drops the LLM sink and a
    # sensitive source + cloud LLM is wrongly allowed.
    snap = {"llm.provider": "anthropic"}  # cloud provider in the snapshot
    with patch.object(rc, "engine_label", return_value=Label(S, CON)):
        d = rc.audit_run_from_snapshot(
            snap, ctx(primary="collection_x"), "collection_x", llm_provider=""
        )
    assert d is not None and not d.allowed
    assert d.reason == "sensitive_to_exposing_inference"


def test_audit_run_from_snapshot_audits_llm_provider_override():
    # The run's per-request model_provider WINS over the snapshot (get_llm uses
    # the explicit arg). A saved local default must not mask a cloud override
    # chosen for this one run — else queue/follow-up/chat runs leak.
    snap = {"llm.provider": "ollama"}  # saved default is local/contained
    with patch.object(rc, "engine_label", return_value=Label(S, CON)):
        d = rc.audit_run_from_snapshot(
            snap,
            ctx(primary="collection_x"),
            "collection_x",
            llm_provider="anthropic",  # per-request cloud override
        )
    assert d is not None and not d.allowed
    assert d.reason == "sensitive_to_exposing_inference"


def test_audit_run_from_snapshot_fails_closed_on_bad_snapshot():
    # Provider resolution runs inside the fail-closed guard: a malformed
    # snapshot refuses the run rather than escaping the helper uncaught.
    d = rc.audit_run_from_snapshot(
        "not-a-dict", ctx(primary="collection_x"), "collection_x"
    )
    assert d is not None and not d.allowed and d.reason == "audit_error"


# ---------------------------------------------------------------------------
# provider endpoint URL fail-up (only ever tightens)
# ---------------------------------------------------------------------------


def test_llm_label_local_custom_endpoint_is_contained():
    # A self-hosted OpenAI-compatible endpoint on a local IP must be CONTAINED
    # (regression: previously mislabeled EXPOSING -> false 400 on local runs).
    snap = {"llm.openai_endpoint.url": "http://127.0.0.1:8000/v1"}
    assert rc.llm_label("openai_endpoint", snap, ctx()) == Label(NS, CON)


def test_llm_label_public_custom_endpoint_is_exposing():
    snap = {"llm.openai_endpoint.url": "http://8.8.8.8:8000/v1"}
    assert rc.llm_label("openai_endpoint", snap, ctx()) == Label(NS, EXP)


# ---------------------------------------------------------------------------
# Real declared label values — guards against a typo in any engine class
# ---------------------------------------------------------------------------


def test_registered_engines_declare_consistent_labels():
    from local_deep_research.security.egress.policy import _get_engine_class
    from local_deep_research.web_search_engines.engine_registry import (
        ENGINE_REGISTRY,
    )

    assert ENGINE_REGISTRY, "engine registry unexpectedly empty"
    checked = 0
    for name in ENGINE_REGISTRY:
        cls = _get_engine_class(name)  # the real class the resolver reads
        if cls is None:
            continue
        sens = getattr(cls, "egress_sensitivity", None)
        exp = getattr(cls, "egress_exposure", None)
        assert sens is not None, f"{name} missing egress_sensitivity"
        assert exp is not None, f"{name} missing egress_exposure"
        if getattr(cls, "is_public", False):
            assert sens is Sensitivity.NON_SENSITIVE, name
            assert exp is Exposure.EXPOSING, name
        elif getattr(cls, "is_local", False):
            assert sens is Sensitivity.SENSITIVE, name
            assert exp is Exposure.CONTAINED, name
        checked += 1
    assert checked > 10, f"only checked {checked} engines"


def test_llm_provider_declared_exposure_matches_resolver():
    # The resolver classifies a provider's endpoint via the enforcing PEP; with
    # no configured URL it falls back to the static cloud/local split. Assert
    # the resolver AND (where the class is resolvable) the declared attribute
    # agree, across every statically-classifiable provider — url-configurable
    # ones (openai_endpoint / anthropic_endpoint) are refined by URL and are
    # covered by the dedicated local/public endpoint tests above.
    from local_deep_research.llm.providers import get_provider_class

    expected = {
        "ollama": Exposure.CONTAINED,
        "lmstudio": Exposure.CONTAINED,
        "llamacpp": Exposure.CONTAINED,
        "openai": Exposure.EXPOSING,
        "anthropic": Exposure.EXPOSING,
        "google": Exposure.EXPOSING,
        "openrouter": Exposure.EXPOSING,
        "deepseek": Exposure.EXPOSING,
        "xai": Exposure.EXPOSING,
        "ionos": Exposure.EXPOSING,
    }
    for name, exp in expected.items():
        assert rc.llm_label(name).exposure is exp, name
        cls = None
        try:
            cls = get_provider_class(name) or get_provider_class(name.upper())
        except Exception:
            cls = None
        if cls is not None:
            assert getattr(cls, "egress_exposure", None) is exp, name


# ---------------------------------------------------------------------------
# Per-destination trust (stage D) — relaxes an exposing sink to contained
# ---------------------------------------------------------------------------


def test_llm_trust_relaxes_cloud_provider_to_contained():
    trusted = {"policy.trusted_inference_providers": ["anthropic"]}
    assert rc.llm_label("anthropic", trusted, ctx()).exposure is CON
    # untrusted cloud stays exposing
    assert rc.llm_label("anthropic", {}, ctx()).exposure is EXP


def test_embeddings_trust_relaxes_cloud_provider_to_contained():
    trusted = {"policy.trusted_inference_providers": ["openai"]}
    assert rc.embeddings_label("openai", trusted, ctx()).exposure is CON


def test_engine_trust_relaxes_exposing_store_to_contained():
    # A sensitive store on a public URL is quadrant 4 (exposing); trusting it
    # moves exposure to contained while keeping its data sensitive.
    trusted = {"policy.trusted_search_engines": ["paperless"]}
    with (
        patch(f"{POLICY}._get_engine_class", return_value=_LocalStore),
        patch(f"{POLICY}._classify_engine_url", return_value=False),
    ):
        assert rc.engine_label("paperless", trusted, ctx()) == Label(S, CON)


def test_trusted_names_handles_json_string_and_garbage():
    assert rc._trusted_names("k", {"k": '["Anthropic", "ollama"]'}) == {
        "anthropic",
        "ollama",
    }
    assert rc._trusted_names("k", {"k": "not json"}) == set()
    assert rc._trusted_names("k", None) == set()


def test_trusted_cloud_llm_allows_sensitive_run_end_to_end():
    # The stage-D payoff: a sensitive source + a TRUSTED cloud LLM is allowed.
    trusted = {"policy.trusted_inference_providers": ["anthropic"]}
    with patch.object(rc, "engine_label", return_value=Label(S, CON)):
        d = rc.audit_run(
            trusted, ctx(), engines=["collection_x"], llm_provider="anthropic"
        )
    assert d is not None and d.allowed


def test_engine_trust_does_not_relax_public_engine():
    # Trusting an inherently-public engine must NOT contain it (no laundering
    # a public search sink past the sensitive->exposing rule).
    class _PublicWithName:
        is_public = True
        egress_sensitivity = NS
        egress_exposure = EXP

    trusted = {"policy.trusted_search_engines": ["searxng"]}
    with patch(f"{POLICY}._get_engine_class", return_value=_PublicWithName):
        assert rc.engine_label("searxng", trusted, ctx()) == Label(NS, EXP)