diff --git a/replicant/web/server.py b/replicant/web/server.py index 91c54e2..8bce469 100644 --- a/replicant/web/server.py +++ b/replicant/web/server.py @@ -624,6 +624,38 @@ def get_catalog() -> dict[str, Any]: "techniques": _technique_json(catalog), } + # Sample lines, cached by what determines them. + # + # The endpoint returns three lines and used to build the entire plan to get + # them, every time the operator selected a technique: REP-004 at high + # intensity is 180,000 events and about 1.6 seconds of pure CPU, discarded + # immediately. It cannot simply build three and stop, because the sample is + # deliberately the first, middle and last event, and the middle of three is + # not the middle of the run. + # + # Caching is exact rather than approximate: the request is built from a fixed + # seed and a fixed intensity, so (technique, intensity, vendor) determines the + # output completely. Bounded because a server is long-lived and the catalog + # times intensities times vendors is 24 * 3 * 3, which fits comfortably. + sample_cache: dict[tuple[str, str, str], list[str]] = {} + SAMPLE_CACHE_MAX = 256 + + def _sample_lines(orch: Orchestrator, request: RunRequest, vendor_id: str) -> list[str]: + key = (request.technique_id, request.intensity, vendor_id) + cached = sample_cache.get(key) + if cached is not None: + return cached + events = list(orch.build_plan(request).events) + if events: + idxs = sorted({0, len(events) // 2, len(events) - 1}) + lines = [orch.render_line(events[i]) for i in idxs] + else: + lines = [] + if len(sample_cache) >= SAMPLE_CACHE_MAX: # pragma: no cover - 216 possible keys + sample_cache.clear() + sample_cache[key] = lines + return lines + @app.get("/api/catalog/{technique_id}/sample", dependencies=[Depends(require_token)]) def technique_sample( technique_id: str, @@ -641,12 +673,7 @@ def technique_sample( seed=settings.default_seed, no_send=True, ) - events = list(orch.build_plan(request).events) - if events: - idxs = sorted({0, len(events) // 2, len(events) - 1}) - lines = [orch.render_line(events[i]) for i in idxs] - else: - lines = [] + lines = _sample_lines(orch, request, _resolve_vendor(vendor)) return { "technique_id": technique.id, "vendor": _resolve_vendor(vendor), diff --git a/tests/test_sample_cache.py b/tests/test_sample_cache.py new file mode 100644 index 0000000..65f80cc --- /dev/null +++ b/tests/test_sample_cache.py @@ -0,0 +1,113 @@ +# Copyright 2026 Imran Hafeez (RZA) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Three sample lines should not cost a whole plan, every time. + +Review finding #3. The sample endpoint fires whenever the operator selects a +technique in the catalog rail, and it built the entire plan to pick three lines +out of it. REP-004 at high intensity is 180,000 events and about 1.6 seconds of +pure CPU, discarded immediately. + +The reviewer's suggested fix, "build three and stop", would change what the +sample shows: it is deliberately the FIRST, MIDDLE and LAST event, and the middle +of three is not the middle of the run. The sample is fully determined by +(technique, intensity, vendor) given a fixed seed, so caching preserves the +output exactly and removes the repeat cost. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +pytest.importorskip("fastapi") + +from fastapi.testclient import TestClient # noqa: E402 + +from replicant.config.settings import Settings # noqa: E402 +from replicant.core.models import load_catalog # noqa: E402 +from replicant.resources import TECHNIQUE_CATALOG # noqa: E402 +from replicant.web.server import create_app # noqa: E402 + +TOKEN = "test-token" +HEADERS = {"x-replicant-token": TOKEN} +CATALOG = load_catalog(TECHNIQUE_CATALOG) + + +@pytest.fixture() +def client(tmp_path: Path) -> TestClient: + app = create_app(CATALOG, Settings(manifest_dir=str(tmp_path)), token=TOKEN) + return TestClient(app, base_url="http://localhost") + + +def _sample(client: TestClient, tid: str = "REP-004", intensity: str = "high") -> dict: + resp = client.get( + f"/api/catalog/{tid}/sample", headers=HEADERS, params={"intensity": intensity} + ) + assert resp.status_code == 200, resp.text + return resp.json() + + +def test_a_repeat_request_does_not_rebuild_the_plan( + client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + first = _sample(client) + + # Any further build would be a cache miss. Counting is used rather than + # timing, which passes or fails on how busy the machine is. + from replicant.core.orchestrator import Orchestrator + + builds = {"n": 0} + original = Orchestrator.build_plan + + def counted(self, request): # type: ignore[no-untyped-def] + builds["n"] += 1 + return original(self, request) + + monkeypatch.setattr(Orchestrator, "build_plan", counted) + second = _sample(client) + + assert builds["n"] == 0, "the plan was rebuilt for an identical sample request" + assert second["lines"] == first["lines"] + + +def test_the_sample_is_still_first_middle_and_last(client: TestClient) -> None: + """The cache must not change what a sample contains.""" + body = _sample(client, "REP-001", "low") + + assert body["lines"] + assert len(body["lines"]) <= 3 + assert all(line.startswith("CEF:0|") for line in body["lines"]) + + +def test_a_different_intensity_is_a_different_sample(client: TestClient) -> None: + """Intensity is part of the key, so it must not serve a stale answer.""" + low = _sample(client, "REP-004", "low") + high = _sample(client, "REP-004", "high") + + assert low["intensity"] == "low" + assert high["intensity"] == "high" + + +def test_a_different_vendor_is_a_different_sample(client: TestClient) -> None: + fortigate = client.get( + "/api/catalog/REP-001/sample", headers=HEADERS, params={"vendor": "fortigate"} + ).json() + checkpoint = client.get( + "/api/catalog/REP-001/sample", headers=HEADERS, params={"vendor": "checkpoint"} + ).json() + + assert fortigate["lines"] != checkpoint["lines"] + assert "Fortinet" in fortigate["lines"][0] + assert "Check Point" in checkpoint["lines"][0]