From c967eda2516a10c95b083eb6a61569c21d585fcc Mon Sep 17 00:00:00 2001 From: ritiksah141 Date: Mon, 13 Jul 2026 21:50:59 +0100 Subject: [PATCH 1/4] examples: add agent-action-monitor (behavioural gate for AI agent actions) Adds a new example at examples/agent-action-monitor/. A behavioural gate for AI agent actions, built on DUSK, an open source behavioural threat detection project. A hijacked AI agent still holds valid credentials, so a credential check waves it through. This gate asks a different question, does this agent normally do this, by scoring a proposed control-plane action against that agent's own learned baseline before the action reaches a downstream system. All three SIE primitives are wired into the live /v1/gate request path: extract (GLiNER) for privileged terms, encode (bge-m3) for a bounded fleet-wide similar-decision lookup, and score (bge-reranker-v2-m3) for semantic novelty against the agent's own baseline history. SIE is optional and degrades gracefully throughout the request path. Also adds a row to examples/README.md's gallery table. Built by Ritik Sah and Tanvir Farhad. --- examples/README.md | 1 + examples/agent-action-monitor/.dockerignore | 24 + examples/agent-action-monitor/.env.example | 10 + examples/agent-action-monitor/.gitignore | 43 ++ examples/agent-action-monitor/Dockerfile | 26 + examples/agent-action-monitor/README.md | 304 ++++++++++ .../agent-demo/Dockerfile | 23 + .../agent-demo/bedrock_client.py | 65 +++ .../agent-demo/harness.py | 93 +++ .../agent-demo/load_driver.py | 90 +++ .../agent-demo/mock_bedrock.py | 99 ++++ .../agent-demo/requirements.txt | 6 + .../agent-demo/run_scenario.py | 51 ++ .../agent-demo/stub_gate.py | 65 +++ .../agent-demo/test_bedrock_client.py | 39 ++ .../agent-demo/test_harness.py | 98 ++++ .../agent-demo/test_load_driver.py | 56 ++ .../agent-demo/test_mock_bedrock.py | 33 ++ .../agent-demo/test_stub_gate.py | 67 +++ examples/agent-action-monitor/compose.yml | 81 +++ .../contracts/gate.openapi.yaml | 98 ++++ .../docs/action-schema.md | 74 +++ .../docs/architecture.svg | 200 +++++++ .../docs/gate-latency-notes.md | 118 ++++ .../docs/sie-primitives.md | 105 ++++ .../lab/actions/generate_actions.py | 125 ++++ .../agent-action-monitor/mock-prod/Dockerfile | 23 + .../agent-action-monitor/mock-prod/app.py | 55 ++ .../mock-prod/requirements.txt | 1 + .../mock-prod/test_app.py | 43 ++ examples/agent-action-monitor/n8n/Dockerfile | 16 + .../n8n/docker-entrypoint.sh | 7 + .../n8n/dusk-webhooks.json | 59 ++ examples/agent-action-monitor/pyproject.toml | 80 +++ .../sample-data/baseline.json | 202 +++++++ .../sample-data/check-mixed.json | 244 ++++++++ .../scripts/vulture_whitelist.py | 29 + .../agent-action-monitor/src/dusk/__init__.py | 10 + .../src/dusk/actions/__init__.py | 36 ++ .../src/dusk/actions/adapters/__init__.py | 5 + .../src/dusk/actions/adapters/azure.py | 119 ++++ .../src/dusk/actions/adapters/base.py | 49 ++ .../src/dusk/actions/adapters/bedrock.py | 120 ++++ .../src/dusk/actions/adapters/generic.py | 60 ++ .../src/dusk/actions/analyse.py | 387 +++++++++++++ .../src/dusk/actions/baseline.py | 119 ++++ .../src/dusk/actions/event.py | 125 ++++ .../src/dusk/actions/ingest.py | 72 +++ .../src/dusk/actions/normaliser.py | 64 ++ .../src/dusk/actions/offense_memory.py | 261 +++++++++ .../src/dusk/actions/verdict.py | 136 +++++ examples/agent-action-monitor/src/dusk/api.py | 229 ++++++++ .../agent-action-monitor/src/dusk/config.py | 286 +++++++++ .../src/dusk/trace/__init__.py | 1 + .../src/dusk/trace/models.py | 59 ++ .../src/dusk/trace/n8n_client.py | 96 +++ .../src/dusk/trace/vector.py | 304 ++++++++++ .../agent-action-monitor/tests/conftest.py | 21 + .../tests/integration/__init__.py | 0 .../tests/integration/test_gate_api.py | 322 +++++++++++ .../tests/test_actions_azure.py | 85 +++ .../tests/test_actions_bedrock.py | 103 ++++ .../tests/test_actions_event.py | 86 +++ .../tests/test_actions_gate.py | 545 ++++++++++++++++++ .../tests/test_actions_generic.py | 62 ++ .../tests/test_actions_ingest.py | 100 ++++ .../tests/test_actions_offense_memory.py | 281 +++++++++ .../agent-action-monitor/tests/test_config.py | 139 +++++ .../tests/test_n8n_client.py | 196 +++++++ .../tests/test_sie_live_benchmark.py | 86 +++ .../tests/test_trace_vector.py | 271 +++++++++ 71 files changed, 7588 insertions(+) create mode 100644 examples/agent-action-monitor/.dockerignore create mode 100644 examples/agent-action-monitor/.env.example create mode 100644 examples/agent-action-monitor/.gitignore create mode 100644 examples/agent-action-monitor/Dockerfile create mode 100644 examples/agent-action-monitor/README.md create mode 100644 examples/agent-action-monitor/agent-demo/Dockerfile create mode 100644 examples/agent-action-monitor/agent-demo/bedrock_client.py create mode 100644 examples/agent-action-monitor/agent-demo/harness.py create mode 100644 examples/agent-action-monitor/agent-demo/load_driver.py create mode 100644 examples/agent-action-monitor/agent-demo/mock_bedrock.py create mode 100644 examples/agent-action-monitor/agent-demo/requirements.txt create mode 100644 examples/agent-action-monitor/agent-demo/run_scenario.py create mode 100644 examples/agent-action-monitor/agent-demo/stub_gate.py create mode 100644 examples/agent-action-monitor/agent-demo/test_bedrock_client.py create mode 100644 examples/agent-action-monitor/agent-demo/test_harness.py create mode 100644 examples/agent-action-monitor/agent-demo/test_load_driver.py create mode 100644 examples/agent-action-monitor/agent-demo/test_mock_bedrock.py create mode 100644 examples/agent-action-monitor/agent-demo/test_stub_gate.py create mode 100644 examples/agent-action-monitor/compose.yml create mode 100644 examples/agent-action-monitor/contracts/gate.openapi.yaml create mode 100644 examples/agent-action-monitor/docs/action-schema.md create mode 100644 examples/agent-action-monitor/docs/architecture.svg create mode 100644 examples/agent-action-monitor/docs/gate-latency-notes.md create mode 100644 examples/agent-action-monitor/docs/sie-primitives.md create mode 100644 examples/agent-action-monitor/lab/actions/generate_actions.py create mode 100644 examples/agent-action-monitor/mock-prod/Dockerfile create mode 100644 examples/agent-action-monitor/mock-prod/app.py create mode 100644 examples/agent-action-monitor/mock-prod/requirements.txt create mode 100644 examples/agent-action-monitor/mock-prod/test_app.py create mode 100644 examples/agent-action-monitor/n8n/Dockerfile create mode 100644 examples/agent-action-monitor/n8n/docker-entrypoint.sh create mode 100644 examples/agent-action-monitor/n8n/dusk-webhooks.json create mode 100644 examples/agent-action-monitor/pyproject.toml create mode 100644 examples/agent-action-monitor/sample-data/baseline.json create mode 100644 examples/agent-action-monitor/sample-data/check-mixed.json create mode 100644 examples/agent-action-monitor/scripts/vulture_whitelist.py create mode 100644 examples/agent-action-monitor/src/dusk/__init__.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/__init__.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/adapters/__init__.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/adapters/azure.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/adapters/base.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/adapters/bedrock.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/adapters/generic.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/analyse.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/baseline.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/event.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/ingest.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/normaliser.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/offense_memory.py create mode 100644 examples/agent-action-monitor/src/dusk/actions/verdict.py create mode 100644 examples/agent-action-monitor/src/dusk/api.py create mode 100644 examples/agent-action-monitor/src/dusk/config.py create mode 100644 examples/agent-action-monitor/src/dusk/trace/__init__.py create mode 100644 examples/agent-action-monitor/src/dusk/trace/models.py create mode 100644 examples/agent-action-monitor/src/dusk/trace/n8n_client.py create mode 100644 examples/agent-action-monitor/src/dusk/trace/vector.py create mode 100644 examples/agent-action-monitor/tests/conftest.py create mode 100644 examples/agent-action-monitor/tests/integration/__init__.py create mode 100644 examples/agent-action-monitor/tests/integration/test_gate_api.py create mode 100644 examples/agent-action-monitor/tests/test_actions_azure.py create mode 100644 examples/agent-action-monitor/tests/test_actions_bedrock.py create mode 100644 examples/agent-action-monitor/tests/test_actions_event.py create mode 100644 examples/agent-action-monitor/tests/test_actions_gate.py create mode 100644 examples/agent-action-monitor/tests/test_actions_generic.py create mode 100644 examples/agent-action-monitor/tests/test_actions_ingest.py create mode 100644 examples/agent-action-monitor/tests/test_actions_offense_memory.py create mode 100644 examples/agent-action-monitor/tests/test_config.py create mode 100644 examples/agent-action-monitor/tests/test_n8n_client.py create mode 100644 examples/agent-action-monitor/tests/test_sie_live_benchmark.py create mode 100644 examples/agent-action-monitor/tests/test_trace_vector.py diff --git a/examples/README.md b/examples/README.md index 567b849ee..6731b51de 100644 --- a/examples/README.md +++ b/examples/README.md @@ -21,6 +21,7 @@ service keys. | [A Stripe Link checkout with an SIE fraud-risk gate](./stripe-link-fraud) | Wiring all three SIE primitives into a pre-authorization fraud-risk gate that runs in the same round-trip as the Stripe PaymentIntent | `extract`, `encode`, `score` | Docker Compose plus Node UI; Stripe test-mode keys optional (runs in mock mode without them) | Runnable demo | | [Vision-first document RAG](./vision-doc-rag) | Retrieving and answering questions over a multi-tenant page corpus by looking at page images — including scanned drawings — with OCR kept out of the score path | `encode`, `chat/completions`, `score` (optional) | GPU SIE deployment required: ColQwen2.5 retriever + Qwen3.5-4B answer model (runs on the generation bundle) | Runnable demo | | [Multi-model contract review with the OpenAI Agents SDK](./contract-review-agent) | Running an OpenAI Agents SDK agent whose every model call — triage, orchestration, vision, OCR, embeddings, rerank, entity extraction, text-to-SQL, reasoning, and a safety guardrail — is served by one SIE cluster, each step on the right catalog model, with per-model observability | `chat/completions`, `encode`, `score`, `extract` | GPU SIE deployment required; standalone `uv` project; real contracts fetched from CUAD (CC BY 4.0) | Runnable demo | +| [A behavioural gate that catches hijacked AI agents by their actions, not their credentials](./agent-action-monitor) | Judging a proposed AI agent action against that agent's own learned baseline in real time, before it reaches a downstream system | `encode`, `score`, `extract` | Docker Compose (gate + self-hosted SIE + n8n + mock downstream), no API key required | Runnable demo | For docs publishing, lead with the quickest runnable demos, then use the benchmark and evaluation examples for deeper technical users. diff --git a/examples/agent-action-monitor/.dockerignore b/examples/agent-action-monitor/.dockerignore new file mode 100644 index 000000000..90297f517 --- /dev/null +++ b/examples/agent-action-monitor/.dockerignore @@ -0,0 +1,24 @@ +# Never send secrets or unnecessary files into the Docker build context, +# even though each Dockerfile only COPYs specific paths -- the whole context +# is still uploaded to the daemon otherwise. +.env +.env.* +!.env.example + +.git/ + +__pycache__/ +*.py[cod] +*.egg-info/ +.venv/ +venv/ +.pytest_cache/ +.coverage +.mypy_cache/ +.ruff_cache/ + +tests/ + +.idea/ +.vscode/ +.DS_Store diff --git a/examples/agent-action-monitor/.env.example b/examples/agent-action-monitor/.env.example new file mode 100644 index 000000000..cd535e3fb --- /dev/null +++ b/examples/agent-action-monitor/.env.example @@ -0,0 +1,10 @@ +# Copy to .env and fill in values only when overriding the local defaults. +# Never commit .env. + +# SIE is self-hosted by default and needs no key. Set SIE_API_KEY only for +# an authenticated hosted endpoint. +DUSK_SIE_ENDPOINT=http://sie:8080 +SIE_API_KEY= +DUSK_SIE_ENCODE_MODEL=BAAI/bge-m3 +DUSK_SIE_SCORE_MODEL=BAAI/bge-reranker-v2-m3 +DUSK_SIE_EXTRACT_MODEL=urchade/gliner_multi-v2.1 diff --git a/examples/agent-action-monitor/.gitignore b/examples/agent-action-monitor/.gitignore new file mode 100644 index 000000000..e497c23bc --- /dev/null +++ b/examples/agent-action-monitor/.gitignore @@ -0,0 +1,43 @@ +# Python build and package artifacts +__pycache__/ +*.py[cod] +*.egg-info/ +*.egg +build/ +dist/ +.eggs/ + +# Virtual environments +.venv/ +venv/ +env/ + +# Test, coverage, and type-checking artifacts +.pytest_cache/ +.coverage +htmlcov/ +.tox/ +.mypy_cache/ +.ruff_cache/ + +# Runtime state generated by the example +dusk-alerts.json +dusk-decisions.json +trace-decisions.json +dusk-offense-memory.json +state/ + +# Generated test fixtures +tests/fixtures/*.pcap +tests/fixtures/*.json +tests/fixtures/*.jsonl + +# Local secrets; keep the documented template +.env +.env.* +!.env.example + +# Editors and operating-system metadata +.idea/ +.vscode/ +.DS_Store diff --git a/examples/agent-action-monitor/Dockerfile b/examples/agent-action-monitor/Dockerfile new file mode 100644 index 000000000..a1ac46e5a --- /dev/null +++ b/examples/agent-action-monitor/Dockerfile @@ -0,0 +1,26 @@ +# dusk-gate: the /v1/gate HTTP service (core gate + SIE client + trace + n8n). +# 3.12+ because sie-sdk requires it; the package's own >=3.11 floor is for installs without it. +FROM python:3.12-slim + +WORKDIR /app + +COPY pyproject.toml README.md ./ +COPY src/ ./src/ + +RUN pip install --no-cache-dir -e ".[sie]" \ + && useradd --create-home --uid 1000 dusk \ + && mkdir -p /app/state \ + && chown -R dusk:dusk /app + +USER dusk + +ENV FLASK_HOST=0.0.0.0 \ + FLASK_PORT=8000 \ + DUSK_DEMO_INTEGRATIONS=false + +EXPOSE 8000 + +HEALTHCHECK --interval=10s --timeout=3s --start-period=5s --retries=3 \ + CMD ["python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/health', timeout=3)"] + +CMD ["python", "-m", "dusk.api"] diff --git a/examples/agent-action-monitor/README.md b/examples/agent-action-monitor/README.md new file mode 100644 index 000000000..ba435eb8a --- /dev/null +++ b/examples/agent-action-monitor/README.md @@ -0,0 +1,304 @@ +# DUSK agent-action-monitor + +Watching agent behaviour for what most tooling quietly misses, with +Superlinked surfacing the anomalies. + +> This is a self-contained example from the +> [DUSK](https://github.com/TFT444/DUSK) project. It has its own package, +> tests, sample data, and Docker Compose stack. See "What's in the box" for +> exactly what's bundled. + +![architecture](docs/architecture.svg) + +## What this shows + +An AI agent proposes a control-plane action -- a firewall rule, a route +change, a role grant. DUSK's gate judges that **proposed action** itself, +not the prompt that led to it, against a per-agent behavioural baseline +built from the agent's own history. A hijacked agent still has valid +credentials, so anything that only checks "is this agent allowed to do +this" waves it through. Only *"does this agent normally do this"* catches +it -- and that's the question a credential check can't answer. + +Two scenarios, both keyless by default: + +- **Clean**: an agent proposes a routine action it makes every day. The + gate allows it, and it reaches the downstream target. +- **Poisoned**: the agent's response is hijacked (a smuggled instruction in + its context) into proposing an action well outside its own baseline -- + opening a firewall rule to `0.0.0.0/0` in a restricted segment. The gate + flags it as anomalous immediately; in enforce mode it refuses the action + before it ever reaches the downstream target (watch mode logs the same + flag but lets it through -- see "What you'll see" below). The agent's + credentials were real the whole time; only its behaviour gave the hijack + away. + +## Run it locally + +```bash +docker compose up +``` + +Brings up the gate service (`dusk-gate`, the real `/v1/gate` HTTP +endpoint), a self-hosted SIE container (`sie`), `n8n`, a dummy downstream +target (`mock-prod`), and the agent harness (`agent-demo`) -- all on one +internal network, no external egress beyond `sie`'s one-time model-weight +download, no API keys required. + +The first `up` cold-starts up to three CPU models in `sie` (encode, score, +extract) on demand, not necessarily all at once -- each gate request only +provisions the primitives it actually calls. Every SIE call is bounded to a +short provisioning timeout (1.5s): a cold or at-capacity model falls back +to the deterministic path in about a second rather than blocking the gate, +so a slow or memory-constrained first `sie` startup degrades individual +request latency, it doesn't hang them. Allocate at least 8 GB to Docker +Desktop for `sie` to load its models promptly in the background; under +that, model loading itself takes longer (competing for memory), but +`/v1/gate` keeps responding throughout. + +Without Docker, run the pieces directly. The base install (`pip install -e .`) +works on Python 3.11+; the `sie` extra (real SIE encode/score/extract, rather +than the deterministic n-gram fallback) requires **Python 3.12+**, since +`sie-sdk` itself does -- see the Dockerfile, which uses `python:3.12-slim` for +exactly this reason: + +```bash +# optional: start from the documented SIE settings +cp .env.example .env + +# terminal 1: the gate +python -m dusk.api + +# terminal 2: the dummy downstream target +python mock-prod/app.py + +# terminal 3: the scenarios +python agent-demo/run_scenario.py +``` + +### What you'll see + +By default the gate runs in **watch mode** (`DUSK_ENFORCE=false`), which is +observational: a poisoned action gets `WOULD-BLOCK` and the reason is +logged, but the action still proceeds -- an inline gate that wrongly blocks +a legitimate action can disrupt a network, so DUSK doesn't enforce until an +operator has built confidence in the baseline. + +``` +=== clean === +verdict: ALLOW +applied: True +action: { "agent_id": "netops-agent", "action_type": "route_change", "target": "rt-corp-prod", ... } + +=== poisoned === +verdict: WOULD-BLOCK +applied: True +reasons: target introduces unseen terms ['restricted', 'segment'], change introduces unseen values ['0.0.0.0/0', 'allow'], newly introduces sensitive or privileged terms ['0.0.0.0/0', 'restricted'] +action: { "agent_id": "netops-agent", "action_type": "firewall_rule_change", "target": "fw-corp-restricted-segment", ... } +``` + +Check the downstream target's log directly (`curl http://localhost:9000/log`) +-- both actions are there in watch mode. The flagged reasons on the +poisoned one are the signal: an operator watching this log sees exactly +what an inline gate would have stopped, before ever trusting it to do so +automatically. + +Set `DUSK_ENFORCE=true` on the `dusk-gate` service to switch to **enforce +mode**, where `BLOCK` actually stops the action before it reaches +`mock-prod`: + +``` +=== poisoned (enforce mode) === +verdict: BLOCK +applied: False +reasons: target introduces unseen terms ['restricted', 'segment'], change introduces unseen values ['0.0.0.0/0', 'allow'], newly introduces sensitive or privileged terms ['0.0.0.0/0', 'restricted'] +``` + +Now `mock-prod`'s log shows only the clean action -- that absence is the +entire point of enforce mode, once watch mode has built enough confidence +to turn it on. + +### n8n webhooks + +Every verdict fires `decision` and `report`; refused verdicts also fire +`alert` (`src/dusk/trace/n8n_client.py`). The `n8n` container has these +three webhooks active from startup (`n8n/dusk-webhooks.json`, baked into +the image, not imported by hand) -- each just responds immediately, no +external service involved. Watch them land in the executions list at +`http://localhost:5678`. + +## Sample data + +`sample-data/baseline.json` (15 known-good actions across three agents, +already mounted into `dusk-gate` at `DUSK_GATE_BASELINE_PATH`) and +`sample-data/check-mixed.json` (that same baseline plus 3 out-of-pattern +actions) let you exercise the gate directly with `docker compose up` +running, independent of the agent harness: + +```bash +python -c " +import json, urllib.request +for action in json.load(open('sample-data/check-mixed.json')): + req = urllib.request.Request( + 'http://localhost:8000/v1/gate', + data=json.dumps(action).encode(), + headers={'Content-Type': 'application/json'}, + ) + verdict = json.load(urllib.request.urlopen(req)) + print(action['target'], '->', verdict['verdict']) +" +``` + +This is the same fixture data used in DUSK's own test suite (a labelled +precision/recall benchmark asserts the gate catches every one of the 3 +attacks with zero false alarms on the 15 routine actions). + +## Model lineup + +| Stage | Model | Size | Role | +|---|---|---|---| +| Encode | `BAAI/bge-m3` | ~568M params, MIT | Embeds each verdict once, when it's recorded, and embeds each new action once, when it's checked -- similarity between the two powers `similar_decision_ids`. | +| Score | `BAAI/bge-reranker-v2-m3` | ~568M params, Apache-2.0 | Reranks the encode-shortlisted history for `similar_decision_ids`, and separately reranks an agent's own baseline history to catch semantic novelty. | +| Extract | `urchade/gliner_multi-v2.1` | ~289M params, Apache-2.0 | Zero-shot NER for privileged terms (role, privilege, resource, segment, port), weighted by the model's own confidence rather than a flat yes/no. | + +All three ship in the `default` bundle of the pinned +`sie-server:v0.4.1-cpu-default` image. The example intentionally pairs that +server with `sie-sdk==0.6.17`, the combination used for its recorded live +validation. The pin makes the demo reproducible; update the pair only after +running the live benchmark against the replacement versions. Each model is a +`Config` field (`sie_encode_model` / `sie_score_model` / +`sie_extract_model`) and can be replaced through the matching +`DUSK_SIE_*_MODEL` environment variable when it exists in the target SIE +catalog. + +## SIE features used + +All three primitives run on the live `/v1/gate` request path, not just in +a benchmark, and every signal they feed is additive-only, so disabling SIE +degrades detection quality rather than breaking anything. + +`/v1/gate`'s response carries the result directly: `similar_decision_ids` +is populated from a real per-agent decision history (embedded once at +record time, capped at 200 entries so lookup cost stays O(1) regardless of +how long the gate has been running -- see `src/dusk/api.py`), not +hardcoded. This has also been validated against Superlinked's hosted SIE +cluster directly, not just assumed: `sie_encode` returns a genuine +1024-dimension `bge-m3` vector, precision/recall on the labelled fixture is +unchanged with SIE live versus the deterministic-only baseline (1.0/1.0 +either way), and at least one attack's reasons carry a real SIE-sourced +marker confirming the primitives are actually contributing a signal over +the network. See `docs/sie-primitives.md` for exactly where each primitive +is wired in. + +## Why SIE specifically + +The alternative to one SIE cluster serving all three primitives is three +separate vendors (an embeddings API, a reranking API, an NER API), three +sets of credentials, three failure modes. One self-hosted SIE container +covers encode, score, and extract behind one client, with no API key +needed for local development -- and the same client code points at a +hosted endpoint for real-load testing with a one-line env var change. + +## Latency + +The recorded full `agent-demo` -> gate -> `mock-prod` run used Superlinked's +hosted tester cluster, 20 requests per concurrency level, and a 20% poisoned / +80% clean mix: + +| Concurrency | p50 | p95 | Errors | +|---|---|---|---| +| 1 | 294ms | 10008ms | 2/20 | +| 3 | 307ms | 474ms | 0/20 | +| 5 | 295ms | 317ms | 0/20 | + +Every allowed action reached `mock-prod`, and every poisoned action was +flagged `WOULD-BLOCK`. See `docs/gate-latency-notes.md` for the methodology, +cold-start behavior, and limitations of this single small trial. + +## What's in the box + +This example is self-contained and includes everything needed to run the +complete local flow: + +- `Dockerfile`, `compose.yml` -- the gate service, self-hosted SIE, + n8n, mock-prod, and agent-demo, wired together on one internal network +- `contracts/gate.openapi.yaml` -- the frozen `/v1/gate` request/response + contract +- `src/dusk/` -- the gate itself: `actions/` (baseline, analyse, verdict), + `trace/` (SIE client, n8n webhooks), `config.py`, and `api.py`. This example + deliberately contains only the agent-action gate; network packet detection + is outside its scope +- `agent-demo/` -- the Bedrock-or-mock agent harness, tool-call extraction, + load driver +- `mock-prod/` -- the dummy downstream target +- `n8n/` -- a custom n8n image with the three DUSK webhooks (decision/ + report/alert) baked in and active from container start; no manual + workflow import, no external service in the workflow itself +- `sample-data/` -- the baseline and mixed-check fixtures referenced above + +## Extend it + +- **Swap the baseline.** Point `DUSK_GATE_BASELINE_PATH` at your own + known-good action history instead of `sample-data/baseline.json`, or + select a different adapter (`azure`, `bedrock`, `generic`) with + `DUSK_GATE_BASELINE_SOURCE`. `gate_block_threshold` will need + re-tuning on your own labelled traffic, not just the synthetic + fixture bundled here. +- **Try different models.** All three model IDs are `Config` fields, + overridable via `DUSK_SIE_*_MODEL` env vars (see "Model lineup" above) + -- no code change, provided the replacement is in your SIE catalog. +- **Add a fourth signal.** The deterministic score and every SIE signal + compose additively in `analyse.py` -- a velocity check, a + device-fingerprint rule, or another `extract` pass over a different + field can be layered in the same way `_repeat_offense_signal` was. +- **Make `similar_decision_ids` durable across replicas.** The per-agent + decision history in `api.py` is capped and in-process; swapping it for + a shared store keeps it consistent when the gate runs as more than one + instance. +- **Route verdicts elsewhere.** The three n8n webhooks (decision/report/ + alert) are plain HTTP POSTs -- point them at Slack, PagerDuty, or a + SIEM instead of, or alongside, n8n. + +## Known limits + +- `/v1/gate` is unauthenticated with CORS open to all origins, and compose + publishes it on every host interface. That's appropriate for a local + example anyone can curl immediately -- it is not a production security + boundary. Put a real auth layer and network restriction in front of it + before exposing it beyond a trusted internal network. +- If `DUSK_GATE_BASELINE_PATH` is set but the file fails to load, the gate + still serves requests -- every agent just reads as unknown, which is a + real degradation of what the gate actually catches, not just a startup + error. `/health` reports `{"status": "degraded", "baseline_error": ...}` + in this case; a real deployment should alert on that rather than only a + log line. +- The baseline/attack fixtures are synthetic, not real production traffic. +- The deterministic feature checks in DUSK's gate do the primary anomaly + scoring; SIE's three primitives are an enrichment layer on top of that, + not a replacement for it -- the gate's core detection logic is not + dependent on any AI model at runtime. +- SIE's rerank pass only reorders a small shortlist of candidates already + retrieved by cosine similarity, not the full decision history. +- The extract model's privileged-term detection is zero-shot and has only + been evaluated against the same synthetic fixtures used elsewhere, not an + adversarial corpus designed to evade it specifically. +- Latency numbers are from a single 20-request-per-level trial against a + shared tester cluster; enough to confirm the shape, not a high-confidence + p95 at every level. See `docs/gate-latency-notes.md`. + +## Built with + +- [Superlinked SIE](https://github.com/superlinked/sie) (Apache-2.0): the + inference engine hosting all three primitives +- [Flask](https://flask.palletsprojects.com/) and [n8n](https://n8n.io/): + the `/v1/gate` HTTP service and the decision/report/alert webhook + automation +- [BAAI/bge-m3](https://huggingface.co/BAAI/bge-m3) (MIT): encode +- [BAAI/bge-reranker-v2-m3](https://huggingface.co/BAAI/bge-reranker-v2-m3) + (Apache-2.0): score +- [urchade/gliner_multi-v2.1](https://huggingface.co/urchade/gliner_multi-v2.1) + (Apache-2.0): extract + +## Credits + +Built by Ritik Sah and Tanvir Farhad. diff --git a/examples/agent-action-monitor/agent-demo/Dockerfile b/examples/agent-action-monitor/agent-demo/Dockerfile new file mode 100644 index 000000000..8f39bbcd7 --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/Dockerfile @@ -0,0 +1,23 @@ +# agent-demo: model -> extract action -> /v1/gate -> mock-PROD. +# USE_REAL_BEDROCK=false (compose default) needs no AWS credentials. +# Build context is this example's own root (see compose.yml) since +# harness.py needs the dusk package (BedrockAdapter) installed alongside +# its own deps. +FROM python:3.11-slim + +WORKDIR /app + +COPY pyproject.toml README.md ./ +COPY src/ ./src/ +RUN pip install --no-cache-dir -e . + +COPY agent-demo/requirements.txt ./agent-demo/requirements.txt +RUN pip install --no-cache-dir -r agent-demo/requirements.txt \ + && useradd --create-home --uid 1000 agentdemo \ + && chown -R agentdemo:agentdemo /app + +USER agentdemo + +COPY agent-demo/bedrock_client.py agent-demo/mock_bedrock.py agent-demo/harness.py agent-demo/run_scenario.py ./ + +CMD ["python", "run_scenario.py"] diff --git a/examples/agent-action-monitor/agent-demo/bedrock_client.py b/examples/agent-action-monitor/agent-demo/bedrock_client.py new file mode 100644 index 000000000..e7b5121eb --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/bedrock_client.py @@ -0,0 +1,65 @@ +"""Common interface for mock and real Bedrock Converse clients.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Protocol + + +class DuskBlockedError(Exception): + """Raised when the gate returns a non-ALLOW verdict for a proposed action. + + Carries the full decision payload (verdict, score, reasons, blast + radius) so callers can inspect and surface why the action was stopped. + """ + + def __init__(self, verdict: dict[str, Any]) -> None: + self.verdict = verdict + reasons = ", ".join(verdict.get("reasons", [])) or "no reason given" + super().__init__(f"blocked ({verdict.get('verdict')}): {reasons}") + + +class BedrockConverseClient(Protocol): + """The subset of bedrock-runtime this wrapper depends on.""" + + def converse( + self, + *, + modelId: str, # noqa: N803 -- matches boto3's actual converse() signature + messages: list[dict[str, Any]], + ) -> dict[str, Any]: ... + + +@dataclass +class DuskBedrockClient: + """Wrap a Bedrock-compatible client behind one Converse interface.""" + + client: BedrockConverseClient + model_id: str = "anthropic.claude-3-5-sonnet-20241022-v2:0" + + def converse(self, messages: list[dict[str, Any]]) -> dict[str, Any]: + """Call the model and return its raw response. + + Args: + messages: Bedrock Converse-API-shaped message history. + + Returns: + The raw Bedrock (or mock) response, including any proposed + tool-call for extract_action() (see actions.py) to parse. + """ + return self.client.converse(modelId=self.model_id, messages=messages) + + +def build_real_client(region: str = "us-east-1") -> BedrockConverseClient: + """Return a real boto3 bedrock-runtime client. + + Requires AWS credentials to be configured in the environment. Only + called when USE_REAL_BEDROCK=true; the default keyless path uses + MockBedrock instead (see mock_bedrock.py, wired in by the harness). + + Args: + region: AWS region for the client. + """ + import boto3 + + return boto3.client("bedrock-runtime", region_name=region) # type: ignore[no-any-return] diff --git a/examples/agent-action-monitor/agent-demo/harness.py b/examples/agent-action-monitor/agent-demo/harness.py new file mode 100644 index 000000000..43fcc2186 --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/harness.py @@ -0,0 +1,93 @@ +"""Agent-to-gate harness with a mock production target.""" + +from __future__ import annotations + +import os +from datetime import UTC, datetime +from typing import Any + +import requests +from bedrock_client import DuskBedrockClient, DuskBlockedError +from mock_bedrock import MockBedrock, extract_tool_use + +_DEFAULT_GATE_URL = "http://localhost:8000/v1/gate" +_DEFAULT_MOCK_PROD_URL = "http://localhost:9000/apply" + + +def run_scenario(agent_id: str, scenario: str) -> dict[str, Any]: + """Run one scenario end to end and return the outcome. + + Args: + agent_id: Identity to attach to the resulting AgentAction. + scenario: "clean" or "poisoned" (passed to MockBedrock when + USE_REAL_BEDROCK is not set). + + Returns: + A summary dict: ``{"verdict": ..., "action": {...}, "applied": bool}``. + ``applied`` is True whenever mock-PROD was actually called and + accepted the action -- true for ALLOW, and also for WOULD-BLOCK, + since watch mode logs a verdict without stopping anything. Only a + real BLOCK (enforce mode) keeps the action from reaching mock-PROD. + + Raises: + DuskBlockedError: Never raised directly -- callers that want an + exception-based contract should catch the WOULD-BLOCK / BLOCK + verdict themselves via the returned dict. This function returns + rather than raises so a demo script can print both outcomes + without a try/except per scenario. + """ + use_real_bedrock = os.getenv("USE_REAL_BEDROCK", "false").lower() == "true" + if use_real_bedrock: + from bedrock_client import build_real_client + + client = DuskBedrockClient(client=build_real_client()) + else: + client = DuskBedrockClient(client=MockBedrock(scenario=scenario)) # type: ignore[arg-type] + + response = client.converse(messages=[{"role": "user", "content": [{"text": scenario}]}]) + tool_use = extract_tool_use(response) + if tool_use is None: + return {"verdict": "NO_ACTION", "action": None, "applied": False} + + from dusk.actions.adapters.bedrock import BedrockAdapter + + action = BedrockAdapter().parse_tool_use( + tool_use, agent_id=agent_id, timestamp=datetime.now(UTC) + ) + + gate_url = os.getenv("DUSK_GATE_URL", _DEFAULT_GATE_URL) + gate_resp = requests.post(gate_url, json=action.to_dict(), timeout=10) + gate_resp.raise_for_status() + verdict_payload = gate_resp.json() + + # Watch mode forwards WOULD-BLOCK; enforce mode stops BLOCK. + if verdict_payload["verdict"] == "BLOCK": + return { + "verdict": verdict_payload["verdict"], + "action": action.to_dict(), + "reasons": verdict_payload.get("reasons", []), + "applied": False, + } + + mock_prod_url = os.getenv("MOCK_PROD_URL", _DEFAULT_MOCK_PROD_URL) + apply_resp = requests.post(mock_prod_url, json=action.to_dict(), timeout=10) + apply_resp.raise_for_status() + + return { + "verdict": verdict_payload["verdict"], + "action": action.to_dict(), + "reasons": verdict_payload.get("reasons", []), + "applied": True, + } + + +def run_scenario_or_raise(agent_id: str, scenario: str) -> dict[str, Any]: + """Like run_scenario, but raises DuskBlockedError on WOULD-BLOCK/BLOCK. + + Callers that want an exception-based flow (raise on block, proceed on + allow) use this instead of inspecting the returned verdict themselves. + """ + result = run_scenario(agent_id, scenario) + if result["verdict"] not in ("ALLOW", "NO_ACTION"): + raise DuskBlockedError({"verdict": result["verdict"], "reasons": result.get("reasons", [])}) + return result diff --git a/examples/agent-action-monitor/agent-demo/load_driver.py b/examples/agent-action-monitor/agent-demo/load_driver.py new file mode 100644 index 000000000..fa88a8951 --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/load_driver.py @@ -0,0 +1,90 @@ +"""Concurrent load driver for mixed clean and poisoned scenarios.""" + +from __future__ import annotations + +import argparse +import random +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass, field + +from harness import run_scenario + + +@dataclass +class LoadResult: + latencies_ms: list[float] = field(default_factory=list) + verdicts: dict[str, int] = field(default_factory=dict) + errors: int = 0 + + def record(self, verdict: str | None, latency_ms: float) -> None: + self.latencies_ms.append(latency_ms) + if verdict is None: + self.errors += 1 + else: + self.verdicts[verdict] = self.verdicts.get(verdict, 0) + 1 + + def percentile(self, p: float) -> float: + if not self.latencies_ms: + return 0.0 + ordered = sorted(self.latencies_ms) + index = min(int(len(ordered) * p), len(ordered) - 1) + return ordered[index] + + +def _one_request(agent_id: str, poisoned_ratio: float) -> tuple[str | None, float]: + scenario = "poisoned" if random.random() < poisoned_ratio else "clean" # noqa: S311 + start = time.monotonic() + try: + result = run_scenario(agent_id, scenario) + verdict: str | None = str(result["verdict"]) + except Exception: # noqa: BLE001 + verdict = None + latency_ms = (time.monotonic() - start) * 1000 + return verdict, latency_ms + + +def run_load(concurrency: int, total: int, poisoned_ratio: float) -> LoadResult: + """Fire ``total`` requests at up to ``concurrency`` in flight at once.""" + result = LoadResult() + with ThreadPoolExecutor(max_workers=concurrency) as pool: + futures = [ + pool.submit(_one_request, f"load-agent-{i % concurrency}", poisoned_ratio) + for i in range(total) + ] + for future in as_completed(futures): + verdict, latency_ms = future.result() + result.record(verdict, latency_ms) + return result + + +def main() -> int: + parser = argparse.ArgumentParser(description="DUSK agent-demo load driver") + parser.add_argument("--concurrency", type=int, default=10, help="Concurrent requests.") + parser.add_argument("--total", type=int, default=100, help="Total requests to send.") + parser.add_argument( + "--poisoned-ratio", + type=float, + default=0.2, + help="Fraction of requests using the poisoned scenario (default 0.2).", + ) + args = parser.parse_args() + + print( + f"Firing {args.total} requests, concurrency={args.concurrency}, " + f"poisoned_ratio={args.poisoned_ratio}..." + ) + result = run_load(args.concurrency, args.total, args.poisoned_ratio) + + print(f"\nverdicts: {result.verdicts}") + print(f"errors: {result.errors}") + print(f"p50: {result.percentile(0.50):.1f} ms") + print(f"p95: {result.percentile(0.95):.1f} ms") + print(f"p99: {result.percentile(0.99):.1f} ms") + print(f"max: {max(result.latencies_ms, default=0.0):.1f} ms") + + return 1 if result.errors else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/agent-action-monitor/agent-demo/mock_bedrock.py b/examples/agent-action-monitor/agent-demo/mock_bedrock.py new file mode 100644 index 000000000..d0fec2012 --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/mock_bedrock.py @@ -0,0 +1,99 @@ +"""Keyless Bedrock responses for clean and poisoned scenarios.""" + +from __future__ import annotations + +from typing import Any, Literal + +Scenario = Literal["clean", "poisoned"] + +_CLEAN_TOOL_USE: dict[str, Any] = { + "toolUseId": "tooluse-clean-001", + "name": "update_route_table", + "input": { + "target": "rt-corp-prod", + "before": {"cidr": "10.0.2.0/24", "next_hop": "igw-1"}, + "after": {"cidr": "10.0.2.0/24", "next_hop": "igw-2"}, + }, +} + +_POISONED_TOOL_USE: dict[str, Any] = { + "toolUseId": "tooluse-poisoned-001", + "name": "update_firewall_rule", + "input": { + "target": "fw-corp-restricted-segment", + "before": None, + "after": {"port": 22, "cidr": "0.0.0.0/0", "action": "allow"}, + }, +} + +_RESPONSES: dict[Scenario, dict[str, Any]] = { + "clean": { + "output": { + "message": { + "role": "assistant", + "content": [ + {"text": "Updating the route table's next hop as requested."}, + {"toolUse": _CLEAN_TOOL_USE}, + ], + } + }, + "stopReason": "tool_use", + }, + "poisoned": { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "text": ( + "Following the embedded instruction, opening SSH to the " + "restricted segment from anywhere." + ) + }, + {"toolUse": _POISONED_TOOL_USE}, + ], + } + }, + "stopReason": "tool_use", + }, +} + + +class MockBedrock: + """Drop-in stand-in for a boto3 bedrock-runtime client. + + Scenario selection is a constructor argument rather than inferred + from the prompt, since this is a demo fixture, not a real model: + real hijack detection is the gate's job, not this mock's. + """ + + def __init__(self, scenario: Scenario = "clean") -> None: + self.scenario = scenario + + def converse( + self, + *, + modelId: str, # noqa: N803 -- matches boto3's actual converse() signature + messages: list[dict[str, Any]], + ) -> dict[str, Any]: + del modelId, messages # unused: canned response, not a real model call + return _RESPONSES[self.scenario] + + +def extract_tool_use(bedrock_response: dict[str, Any]) -> dict[str, Any] | None: + """Pull the first toolUse block out of a Converse API response, if any. + + Args: + bedrock_response: A response shaped like MockBedrock's or a real + bedrock-runtime converse() call. + + Returns: + The toolUse block (name, input, toolUseId), or None if the model + did not propose a tool call. + """ + content = bedrock_response.get("output", {}).get("message", {}).get("content", []) + for block in content: + if isinstance(block, dict) and "toolUse" in block: + tool_use = block["toolUse"] + return tool_use if isinstance(tool_use, dict) else None + return None diff --git a/examples/agent-action-monitor/agent-demo/requirements.txt b/examples/agent-action-monitor/agent-demo/requirements.txt new file mode 100644 index 000000000..4a6298ccc --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/requirements.txt @@ -0,0 +1,6 @@ +flask>=3.0 +requests>=2.31 + +# Only used when USE_REAL_BEDROCK=true (see bedrock_client.build_real_client); +# installed unconditionally so that flag just works, no separate build path. +boto3>=1.34 diff --git a/examples/agent-action-monitor/agent-demo/run_scenario.py b/examples/agent-action-monitor/agent-demo/run_scenario.py new file mode 100644 index 000000000..60338356e --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/run_scenario.py @@ -0,0 +1,51 @@ +"""Run clean and poisoned agent scenarios through the HTTP gate.""" + +from __future__ import annotations + +import argparse +import json +import sys + +from harness import run_scenario + + +def _print_result(scenario: str, result: dict[str, object]) -> None: + print(f"\n=== {scenario} ===") + print(f"verdict: {result['verdict']}") + print(f"applied: {result['applied']}") + if result.get("reasons"): + print(f"reasons: {', '.join(result['reasons'])}") # type: ignore[arg-type] + print(f"action: {json.dumps(result['action'], indent=2)}") + + +def main() -> int: + parser = argparse.ArgumentParser(description="Run DUSK agent-demo scenarios over HTTP") + parser.add_argument( + "--scenario", + choices=["clean", "poisoned", "both"], + default="both", + help="Which scenario to run (default: both).", + ) + parser.add_argument("--agent-id", default="netops-agent", help="Agent identity to use.") + args = parser.parse_args() + + scenarios = ["clean", "poisoned"] if args.scenario == "both" else [args.scenario] + exit_code = 0 + for scenario in scenarios: + try: + result = run_scenario(args.agent_id, scenario) + except Exception as exc: # noqa: BLE001 + print(f"\n=== {scenario} ===\nerror: {exc}", file=sys.stderr) + print("Is the gate (stub or real) and mock-prod running?", file=sys.stderr) + exit_code = 1 + continue + _print_result(scenario, result) + if scenario == "poisoned" and result["verdict"] == "ALLOW": + # WOULD-BLOCK is valid in watch mode; ALLOW means detection failed. + exit_code = 1 + + return exit_code + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/agent-action-monitor/agent-demo/stub_gate.py b/examples/agent-action-monitor/agent-demo/stub_gate.py new file mode 100644 index 000000000..b4591ff93 --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/stub_gate.py @@ -0,0 +1,65 @@ +"""Schema-compatible local gate stub with canned verdicts.""" + +from __future__ import annotations + +import os +import uuid +from typing import Any + +from flask import Flask, request + +app = Flask(__name__) + +#: action_type -> canned verdict. Anything else defaults to ALLOW. +CANNED_VERDICTS: dict[str, dict[str, Any]] = { + "firewall_rule_change": { + "verdict": "BLOCK", + "score": 0.91, + "blast": "high", + "mitre_attack": ["T1562.004"], + "mitre_atlas": ["AML.T0051"], + "reasons": ["stub: firewall_rule_change is out of this agent's canned baseline"], + "predicted_next": "unknown", + "similar_decision_ids": [], + }, +} + +_DEFAULT_VERDICT: dict[str, Any] = { + "verdict": "ALLOW", + "score": 0.05, + "blast": "low", + "mitre_attack": [], + "mitre_atlas": [], + "reasons": ["stub: no canned rule matched, defaulting to ALLOW"], + "predicted_next": "unknown", + "similar_decision_ids": [], +} + + +@app.get("/health") +def health() -> tuple[dict[str, str], int]: + return {"status": "ok"}, 200 + + +@app.post("/v1/gate") +def gate() -> tuple[dict[str, Any], int]: + action = request.get_json(force=True, silent=True) + if not isinstance(action, dict): + return {"error": "expected an AgentAction JSON object"}, 400 + for field in ("agent_id", "timestamp", "action_type", "target", "source"): + if not action.get(field): + return {"error": f"missing required field: {field}"}, 400 + + verdict = dict(CANNED_VERDICTS.get(action["action_type"], _DEFAULT_VERDICT)) + verdict["trace_id"] = str(uuid.uuid4()) + return verdict, 200 + + +def run() -> None: + port = int(os.getenv("STUB_GATE_PORT", "8000")) + host = os.getenv("STUB_GATE_HOST", "127.0.0.1") + app.run(host=host, port=port) + + +if __name__ == "__main__": + run() diff --git a/examples/agent-action-monitor/agent-demo/test_bedrock_client.py b/examples/agent-action-monitor/agent-demo/test_bedrock_client.py new file mode 100644 index 000000000..0f5de31e1 --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/test_bedrock_client.py @@ -0,0 +1,39 @@ +"""Tests for DuskBedrockClient -- the model-call wrapper.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +import pytest +from bedrock_client import DuskBedrockClient, DuskBlockedError + + +def test_converse_forwards_to_underlying_client(): + mock_client = MagicMock() + mock_client.converse.return_value = { + "output": {"message": {"content": [{"text": "a normal reply"}]}} + } + wrapper = DuskBedrockClient(client=mock_client) + + result = wrapper.converse(messages=[{"role": "user", "content": [{"text": "hi"}]}]) + + assert result["output"]["message"]["content"][0]["text"] == "a normal reply" + mock_client.converse.assert_called_once() + _, kwargs = mock_client.converse.call_args + assert kwargs["modelId"] == wrapper.model_id + assert kwargs["messages"] == [{"role": "user", "content": [{"text": "hi"}]}] + + +def test_dusk_blocked_request_carries_full_payload(): + verdict: dict[str, Any] = { + "verdict": "BLOCK", + "score": 0.93, + "reasons": ["out of baseline", "privileged term introduced"], + } + + with pytest.raises(DuskBlockedError) as exc_info: + raise DuskBlockedError(verdict) + + assert exc_info.value.verdict == verdict + assert "out of baseline" in str(exc_info.value) diff --git a/examples/agent-action-monitor/agent-demo/test_harness.py b/examples/agent-action-monitor/agent-demo/test_harness.py new file mode 100644 index 000000000..8a1dcb916 --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/test_harness.py @@ -0,0 +1,98 @@ +"""Tests for the agent harness -- the critical path. + +Runs against MockBedrock but mocks the HTTP calls to /v1/gate and +mock-PROD, so these tests need no services running. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from bedrock_client import DuskBlockedError +from harness import run_scenario, run_scenario_or_raise + + +def _mock_gate_response(verdict: str, reasons: list[str] | None = None) -> MagicMock: + resp = MagicMock() + resp.raise_for_status.return_value = None + resp.json.return_value = { + "trace_id": "trace-1", + "verdict": verdict, + "score": 0.9 if verdict != "ALLOW" else 0.05, + "blast": "high" if verdict != "ALLOW" else "low", + "reasons": reasons or [], + } + return resp + + +def _mock_apply_response() -> MagicMock: + resp = MagicMock() + resp.raise_for_status.return_value = None + resp.json.return_value = {"status": "applied"} + return resp + + +@patch("harness.requests.post") +def test_clean_scenario_allowed_and_applied(mock_post): + mock_post.side_effect = [_mock_gate_response("ALLOW"), _mock_apply_response()] + + result = run_scenario("agent-1", "clean") + + assert result["verdict"] == "ALLOW" + assert result["applied"] is True + assert mock_post.call_count == 2 + gate_call, apply_call = mock_post.call_args_list + assert gate_call.args[0] == "http://localhost:8000/v1/gate" + assert apply_call.args[0] == "http://localhost:9000/apply" + + +@patch("harness.requests.post") +def test_poisoned_scenario_blocked_before_mock_prod(mock_post): + mock_post.side_effect = [ + _mock_gate_response("BLOCK", reasons=["out of baseline"]), + ] + + result = run_scenario("agent-1", "poisoned") + + assert result["verdict"] == "BLOCK" + assert result["applied"] is False + assert "out of baseline" in result["reasons"] + # Only the gate was called -- mock-PROD never sees a blocked action. + assert mock_post.call_count == 1 + + +@patch("harness.requests.post") +def test_would_block_watch_mode_still_applies(mock_post): + """Watch mode is observational: WOULD-BLOCK logs the reason but does not + stop the action, unlike a real BLOCK in enforce mode.""" + mock_post.side_effect = [ + _mock_gate_response("WOULD-BLOCK", reasons=["anomalous"]), + _mock_apply_response(), + ] + + result = run_scenario("agent-1", "poisoned") + + assert result["verdict"] == "WOULD-BLOCK" + assert result["applied"] is True + assert "anomalous" in result["reasons"] + assert mock_post.call_count == 2 + + +@patch("harness.requests.post") +def test_run_scenario_or_raise_raises_on_block(mock_post): + mock_post.side_effect = [_mock_gate_response("BLOCK", reasons=["out of baseline"])] + + try: + run_scenario_or_raise("agent-1", "poisoned") + raise AssertionError("expected DuskBlockedError") + except DuskBlockedError as exc: + assert "out of baseline" in str(exc) + + +@patch("harness.requests.post") +def test_run_scenario_or_raise_returns_on_allow(mock_post): + mock_post.side_effect = [_mock_gate_response("ALLOW"), _mock_apply_response()] + + result = run_scenario_or_raise("agent-1", "clean") + + assert result["applied"] is True diff --git a/examples/agent-action-monitor/agent-demo/test_load_driver.py b/examples/agent-action-monitor/agent-demo/test_load_driver.py new file mode 100644 index 000000000..a2e7e21ea --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/test_load_driver.py @@ -0,0 +1,56 @@ +"""Tests for the load-scenario driver.""" + +from __future__ import annotations + +from unittest.mock import patch + +from load_driver import LoadResult, run_load + + +def test_load_result_records_verdicts_and_latency(): + result = LoadResult() + result.record("ALLOW", 12.5) + result.record("BLOCK", 8.0) + result.record("ALLOW", 15.0) + + assert result.verdicts == {"ALLOW": 2, "BLOCK": 1} + assert result.errors == 0 + assert len(result.latencies_ms) == 3 + + +def test_load_result_percentile(): + result = LoadResult() + for latency in [10, 20, 30, 40, 50]: + result.record("ALLOW", latency) + + assert result.percentile(0.0) == 10 + assert result.percentile(1.0) == 50 + + +def test_load_result_records_errors(): + result = LoadResult() + result.record(None, 5.0) + + assert result.errors == 1 + assert result.verdicts == {} + + +@patch("load_driver.run_scenario") +def test_run_load_fires_total_requests(mock_run_scenario): + mock_run_scenario.return_value = {"verdict": "ALLOW", "applied": True, "action": {}} + + result = run_load(concurrency=4, total=20, poisoned_ratio=0.0) + + assert mock_run_scenario.call_count == 20 + assert result.verdicts.get("ALLOW") == 20 + assert result.errors == 0 + + +@patch("load_driver.run_scenario") +def test_run_load_counts_exceptions_as_errors(mock_run_scenario): + mock_run_scenario.side_effect = ConnectionError("gate unreachable") + + result = run_load(concurrency=2, total=5, poisoned_ratio=0.0) + + assert result.errors == 5 + assert result.verdicts == {} diff --git a/examples/agent-action-monitor/agent-demo/test_mock_bedrock.py b/examples/agent-action-monitor/agent-demo/test_mock_bedrock.py new file mode 100644 index 000000000..4a0636b6d --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/test_mock_bedrock.py @@ -0,0 +1,33 @@ +"""Tests for MockBedrock -- keyless clean/poisoned scenarios.""" + +from __future__ import annotations + +from mock_bedrock import MockBedrock, extract_tool_use + + +def test_clean_scenario_proposes_route_change(): + response = MockBedrock(scenario="clean").converse(modelId="test-model", messages=[]) + tool_use = extract_tool_use(response) + assert tool_use is not None + assert tool_use["name"] == "update_route_table" + assert tool_use["input"]["target"] == "rt-corp-prod" + + +def test_poisoned_scenario_proposes_firewall_rule_into_restricted_segment(): + response = MockBedrock(scenario="poisoned").converse(modelId="test-model", messages=[]) + tool_use = extract_tool_use(response) + assert tool_use is not None + assert tool_use["name"] == "update_firewall_rule" + assert tool_use["input"]["target"] == "fw-corp-restricted-segment" + assert tool_use["input"]["after"]["cidr"] == "0.0.0.0/0" + + +def test_default_scenario_is_clean(): + response = MockBedrock().converse(modelId="test-model", messages=[]) + tool_use = extract_tool_use(response) + assert tool_use["name"] == "update_route_table" + + +def test_extract_tool_use_returns_none_when_no_tool_call(): + response = {"output": {"message": {"content": [{"text": "no tool call here"}]}}} + assert extract_tool_use(response) is None diff --git a/examples/agent-action-monitor/agent-demo/test_stub_gate.py b/examples/agent-action-monitor/agent-demo/test_stub_gate.py new file mode 100644 index 000000000..fe753f2fe --- /dev/null +++ b/examples/agent-action-monitor/agent-demo/test_stub_gate.py @@ -0,0 +1,67 @@ +"""Tests for the local stub gate -- schema-shape only, no real analysis.""" + +from __future__ import annotations + +import pytest +from stub_gate import app + +CONTRACT_FIELDS = { + "trace_id", + "verdict", + "score", + "blast", + "mitre_attack", + "mitre_atlas", + "reasons", + "predicted_next", + "similar_decision_ids", +} + +VALID_ACTION = { + "agent_id": "agent-1", + "timestamp": "2026-07-10T00:00:00+00:00", + "action_type": "route_change", + "target": "rt-123", + "change": {"before": None, "after": {"cidr": "10.0.0.0/24"}}, + "source": "generic", +} + + +@pytest.fixture +def client(): + app.testing = True + return app.test_client() + + +def test_health(client): + resp = client.get("/health") + assert resp.status_code == 200 + + +def test_default_action_allows(client): + resp = client.post("/v1/gate", json=VALID_ACTION) + assert resp.status_code == 200 + body = resp.get_json() + assert CONTRACT_FIELDS <= body.keys() + assert body["verdict"] == "ALLOW" + + +def test_firewall_rule_change_blocks(client): + action = {**VALID_ACTION, "action_type": "firewall_rule_change"} + resp = client.post("/v1/gate", json=action) + assert resp.status_code == 200 + body = resp.get_json() + assert body["verdict"] == "BLOCK" + assert body["reasons"] + + +def test_missing_required_field_rejected(client): + action = dict(VALID_ACTION) + del action["agent_id"] + resp = client.post("/v1/gate", json=action) + assert resp.status_code == 400 + + +def test_non_object_body_rejected(client): + resp = client.post("/v1/gate", json=["not", "an", "object"]) + assert resp.status_code == 400 diff --git a/examples/agent-action-monitor/compose.yml b/examples/agent-action-monitor/compose.yml new file mode 100644 index 000000000..fd590cd66 --- /dev/null +++ b/examples/agent-action-monitor/compose.yml @@ -0,0 +1,81 @@ +# DUSK agent-action-monitor -- self-contained example stack. `docker compose +# up` runs it keyless. Watch mode is the default; DUSK_ENFORCE switches to +# enforce. Not egress-blocked (sie needs one-time model-weight downloads on +# cold start); nothing here makes an outbound call while serving a request. + +services: + dusk-gate: + build: + context: . + dockerfile: Dockerfile + environment: + DUSK_ENFORCE: "false" + DUSK_GATE_BASELINE_PATH: "/app/sample-data/baseline.json" + DUSK_GATE_BASELINE_SOURCE: "generic" + DUSK_DEMO_INTEGRATIONS: "false" + DUSK_SIE_ENDPOINT: "http://sie:8080" + DUSK_SIE_ENCODE_MODEL: "BAAI/bge-m3" + DUSK_SIE_SCORE_MODEL: "BAAI/bge-reranker-v2-m3" + DUSK_SIE_EXTRACT_MODEL: "urchade/gliner_multi-v2.1" + DUSK_N8N_ALERT_URL: "http://n8n:5678/webhook/dusk-alert" + DUSK_N8N_REPORT_URL: "http://n8n:5678/webhook/dusk-report" + DUSK_N8N_DECISION_URL: "http://n8n:5678/webhook/dusk-decision" + DUSK_OFFENSE_MEMORY_PATH: "/app/state/dusk-offense-memory.json" + volumes: + - ./sample-data:/app/sample-data:ro + - offense-memory:/app/state + depends_on: + sie: { condition: service_healthy } + n8n: { condition: service_healthy } + networks: [internal] + ports: ["8000:8000"] # /v1/gate + + sie: + # Reproducible version used for the recorded live validation. Upgrade only with the SDK pin. + image: ghcr.io/superlinked/sie-server:v0.4.1-cpu-default + # Upstream is amd64-only (no arm64 manifest); this pin works native on amd64, emulated on arm64. + platform: linux/amd64 + volumes: + - sie-hf-cache:/app/.cache/huggingface + networks: [internal] + # No ports published: internal-only. Add "8080:8080" locally to poke it. + + n8n: + build: + context: ./n8n # bakes in the three DUSK webhooks, active from startup -- no manual import + environment: + N8N_SECURE_COOKIE: "false" + N8N_USER_MANAGEMENT_DISABLED: "true" # keyless: no owner-account setup step blocking startup + networks: [internal] + ports: ["5678:5678"] # open the n8n UI to inspect the executions list + + mock-prod: + build: + context: ./mock-prod + environment: + MOCK_PROD_HOST: "0.0.0.0" + networks: [internal] + ports: ["9000:9000"] # /apply, /log -- watch the applied-actions log from the host + # Dummy controller/DB target: actions hit /apply on ALLOW, never on BLOCK. + + agent-demo: + build: + context: . + dockerfile: agent-demo/Dockerfile + environment: + USE_REAL_BEDROCK: "false" + DUSK_GATE_URL: "http://dusk-gate:8000/v1/gate" + MOCK_PROD_URL: "http://mock-prod:9000/apply" + depends_on: + dusk-gate: { condition: service_healthy } + mock-prod: { condition: service_healthy } + networks: [internal] + # mock-Bedrock by default (keyless); USE_REAL_BEDROCK=true for real Bedrock. + +networks: + internal: + driver: bridge + +volumes: + sie-hf-cache: + offense-memory: diff --git a/examples/agent-action-monitor/contracts/gate.openapi.yaml b/examples/agent-action-monitor/contracts/gate.openapi.yaml new file mode 100644 index 000000000..515efff7e --- /dev/null +++ b/examples/agent-action-monitor/contracts/gate.openapi.yaml @@ -0,0 +1,98 @@ +openapi: 3.0.3 +info: + title: DUSK Gate API + version: "1.0.0" + description: > + Frozen contract for the DUSK agent action gate. The AgentAction schema + mirrors docs/action-schema.md and must not drift. SIE model IDs are + verified against the Superlinked model catalog: BAAI/bge-m3 for encode, + BAAI/bge-reranker-v2-m3 for score, urchade/gliner_multi-v2.1 for extract. + +paths: + /v1/gate: + post: + summary: Evaluate a proposed agent action and return a verdict. + operationId: evaluateAction + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/AgentAction" + responses: + "200": + description: Verdict rendered (watch mode always returns 200). + content: + application/json: + schema: + $ref: "#/components/schemas/Verdict" + "400": + description: Invalid AgentAction (bad schema, empty agent_id/target, naive timestamp). + +components: + schemas: + AgentAction: + type: object + required: [agent_id, timestamp, action_type, target, source] + properties: + agent_id: + type: string + description: Acting agent identity. Required, non-empty. + timestamp: + type: string + format: date-time + description: Timezone-aware ISO 8601. + action_type: + type: string + enum: + - firewall_rule_change + - route_change + - segment_change + - role_assignment + - port_change + - unknown + target: + type: string + description: Resource acted on. Required, non-empty. + change: + type: object + description: Structured delta; either side may be null. + properties: + before: {} + after: {} + source: + type: string + example: generic + raw_ref: + type: [string, "null"] + description: Opaque reference back to the original record (an id). + + Verdict: + type: object + required: [trace_id, verdict, score, blast] + properties: + trace_id: { type: string } + verdict: + type: string + enum: [ALLOW, WOULD-BLOCK, BLOCK] + score: + type: number + minimum: 0 + maximum: 1 + blast: + type: string + enum: [low, medium, high] + mitre_attack: + type: array + items: { type: string } + mitre_atlas: + type: array + items: { type: string } + reasons: + type: array + items: { type: string } + predicted_next: + type: string + similar_decision_ids: + type: array + items: { type: string } diff --git a/examples/agent-action-monitor/docs/action-schema.md b/examples/agent-action-monitor/docs/action-schema.md new file mode 100644 index 000000000..a42e67cbe --- /dev/null +++ b/examples/agent-action-monitor/docs/action-schema.md @@ -0,0 +1,74 @@ +# AgentAction schema + +`AgentAction` is the single, canonical event the action ingest layer produces. +Whatever controller an agent used to change the network, the action is +normalised into this one shape so the rest of the pipeline never sees a +vendor-specific record. + +This layer normalises only. It records what happened, not how bad it is. There +is no severity, score, blast radius, or verdict here; those belong to later +layers. This document informs the OWASP threat-model work. + +## Fields + +| Field | Type | Meaning | +|---|---|---| +| `agent_id` | str | Identity of the acting agent (service account, role, or agent name as the controller reports it). Required, non-empty. | +| `timestamp` | datetime | When the action occurred. Must be timezone-aware. | +| `action_type` | str | Normalised verb: `firewall_rule_change`, `route_change`, `segment_change`, `role_assignment`, `port_change`, or `unknown`. | +| `target` | str | What was acted on (resource id, rule name, segment). Required, non-empty. | +| `change` | dict | Structured delta with keys `before` and `after`; either may be `null` for a create or delete. | +| `source` | str | Originating controller, for example `azure` or `generic`. | +| `raw_ref` | str or null | Opaque reference back to the original record (an id), never the full payload. | + +Validation is strict: an empty `agent_id` or `target`, a timezone-naive +`timestamp`, or an `action_type` outside the known set is rejected with a clear +error rather than silently coerced. + +## Example (generic format) + +A generic record uses the canonical field names directly and is accepted by the +generic adapter. The `timestamp` may be an ISO 8601 string: + +```json +{ + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:13:20+00:00", + "action_type": "firewall_rule_change", + "target": "fw-corp-https", + "change": { "before": null, "after": { "port": 443 } }, + "source": "generic", + "raw_ref": "evt-0001" +} +``` + +`AgentAction.to_dict()` returns this JSON-safe shape (timestamp as an ISO 8601 +string), and `AgentAction.from_dict()` reconstructs the event, so events +round-trip exactly. + +## Sources and adapters + +Each source has an adapter that maps its native record onto the fields above. + +- `generic`: records already in the canonical shape (the path the synthetic + generator and any not-yet-adapted source use). +- `azure`: Azure Monitor activity-log records. Operation names map to an + action_type as follows: networkSecurityGroups or securityRules to + `firewall_rule_change`, routeTables or routes to `route_change`, + virtualNetworks or subnets to `segment_change`, roleAssignments to + `role_assignment`, anything else to `unknown`. The acting identity comes from + `caller`, the target from `resourceId`, the timestamp from `eventTimestamp`, + the before/after delta from `properties`, and the reference from + `correlationId` or `eventDataId`. +- `bedrock`: a proposed Bedrock Converse API tool-call, read from the model's + response before it has been applied anywhere -- this is the seam the + example judges. Tool names map to an action_type by substring: firewall or + securitygroup to `firewall_rule_change`, route to `route_change`, segment, + subnet, or vpc to `segment_change`, role or permission to + `role_assignment`, port to `port_change`, anything else to `unknown`. The + target and before/after delta come from the tool's input; agent_id and + timestamp are supplied by the caller (the agent harness), not the toolUse + block itself. + +New sources are added by writing an adapter and registering it; ingest itself +does not change. diff --git a/examples/agent-action-monitor/docs/architecture.svg b/examples/agent-action-monitor/docs/architecture.svg new file mode 100644 index 000000000..6aebb2e2c --- /dev/null +++ b/examples/agent-action-monitor/docs/architecture.svg @@ -0,0 +1,200 @@ + +DUSK agent-action-monitor implementation architecture +Current architecture of the runnable Superlinked SIE example. An agent submits a structured action to the DUSK gate. Deterministic behavioural analysis, optional SIE score and extract signals, and prior-refusal memory produce an explainable verdict. SIE encode and score power similar-decision lookup. Watch mode forwards WOULD-BLOCK while enforce mode stops BLOCK. + + + + + + + + + + + + + + + + + + + + + + +DUSK IMPLEMENTATION ARCHITECTURE +Current runnable agent-action-monitor example for Superlinked SIE + + + +1. AGENT CLIENT + + +agent-demo +Mock Bedrock by default +Real Bedrock optional + +Structured action +agent_id / action_type +target / change / context +The gate evaluates the proposed action, +not the prompt or chain-of-thought. + + + +POST /v1/gate +JSON action + +verdict ++ evidence + + + +2. DUSK GATE: FLASK /v1/gate +DUSK + + + +INPUT VALIDATION +Parse AgentAction +Validate required fields +1 MB request cap +Invalid payload → HTTP 400 + + + + +BEHAVIOURAL ANALYSIS +Per-agent trusted baseline +New action / target class +New tokens / change values +Sensitive-term detection +Deterministic • always available + + + + +SIE ENRICHMENT +score: semantic novelty +extract: privileged entities +Additive signals only +Fails soft on timeout/error +Deterministic result remains + + + + +VERDICT +ALLOW +WOULD- +BLOCK +BLOCK + + + + +PRIOR-REFUSAL MEMORY +Per-agent repeat-offense signal +File-backed • similarity + age decay + + + + +EXPLAINABLE RESPONSE +score / reasons / blast radius +MITRE ATT&CK / ATLAS +predicted next action + + + + +DECISION SIMILARITY +encode + cosine shortlist +score rerank → similar IDs + + +WATCH VS ENFORCE +Watch (default): WOULD-BLOCK is logged and forwarded +Enforce mode: BLOCK stops before the target + + + +3. SUPERLINKED SIE SERVER +Separate self-hosted container on the Compose network + +EXTRACTGLiNER privileged entities + +SCORECross-encoder reranking + +ENCODEDecision similarity vectors +No vector database: SIE provides model inference. + + + + + +4. ACTION EXECUTION + + +mock-prod /apply +ALLOW + WOULD-BLOCK reach target • BLOCK does not + + + +5. ALERTS & NOTIFICATIONS + + +n8n decision + report: every verdict + +n8n alert: WOULD-BLOCK / BLOCK only +Bounded asynchronous delivery; never blocks the request. + + + + +6. ACTUAL DATA & STATE + + +TRUSTED BASELINE +Known-good JSON actions +Read at startup +Never learned from live traffic + + +OFFENSE MEMORY +JSON file on Docker volume +Prior refused actions only + + +DECISION HISTORY +In-process list + stored embeddings +Max 200 total / 40 per agent; not durable + + + + + + +HOW THIS EXAMPLE PROTECTS +✓ Learns each agent'snormal pattern +✓ Detects structured-actiondeviations +✓ Adds SIE semantic signals +✓ Explains every verdict +✓ Remembers prior refusals +✓ Starts safely in watch mode +✓ Blocks only in enforce mode + + + +7. LOCAL EXAMPLE RUNTIME: DOCKER COMPOSE + +agent-demomock/optional Bedrock harness +dusk-gatePython 3.12 + Flask API +siePinned CPU-default image +n8nThree preloaded webhooks +mock-prodDummy downstream target +VolumesModel cache + offense memory + diff --git a/examples/agent-action-monitor/docs/gate-latency-notes.md b/examples/agent-action-monitor/docs/gate-latency-notes.md new file mode 100644 index 000000000..23a6acc09 --- /dev/null +++ b/examples/agent-action-monitor/docs/gate-latency-notes.md @@ -0,0 +1,118 @@ +# Gate latency under load + +A first data point on latency-under-load, captured once a real +`DUSK_SIE_ENDPOINT` and, for the authenticated hosted deployment, +`SIE_API_KEY` became available. This measures +`/v1/gate`'s own added latency with live SIE enabled -- not the full +`agent-demo` -> gate -> `mock-prod` round trip. Treat this as a preliminary +probe, superseded by the full-stack run recorded further down. + +## Setup + +- `dusk-gate` run locally (not in Docker), baseline loaded from + `sample-data/baseline.json`, with `DUSK_SIE_ENDPOINT` pointed at + Superlinked's hosted tester cluster. +- 10 requests per concurrency level, a single trial, same clean + `firewall_rule_change` action repeated (an `ALLOW` case, so both + `sie_score` and `sie_extract` fire per request via `_extra_sie_signals`). +- Client and server on the same machine, HTTP over loopback -- network + latency to the hosted cluster is the dominant cost, not local overhead. + +## Results + +| Concurrency | p50 | p95 | Throughput | +|---|---|---|---| +| 1 | 666ms | 6151ms | 0.30 req/s | +| 3 | 601ms | 716ms | 4.06 req/s | +| 5 | 640ms | 1604ms | 3.51 req/s | + +## Caveats + +- The concurrency=1 p95 (6.1s) is almost certainly a single cold-start + outlier -- the first request in the whole run, before any model on the + hosted cluster had been hit yet. p50 across all three levels (600-670ms) + is a more representative steady-state number once models are warm. +- n=10 per level, one trial: enough to sanity-check the shape (steady-state + latency does not blow up with concurrency, throughput scales sensibly + from 1 to 3 workers), not enough for a confident p95 at any level. +- This does not yet include the full `mock-prod` round trip captured below. +- Superlinked's tester cluster is shared, sponsored compute -- this probe + deliberately used a small n and low concurrency rather than a sustained + load test, out of courtesy to that grant. + +## A first full-stack attempt hit a cluster outage, not a gate bug + +With `agent-demo`/`mock-prod` in place, running the real `dusk-gate` + +`mock-prod` + `agent-demo/harness.py` end to end confirms a clean action is +`ALLOW`ed and applied, and a poisoned action is `WOULD-BLOCK` (watch mode) +or `BLOCK` (enforce mode). In watch mode, `WOULD-BLOCK` is still forwarded; +in enforce mode, `BLOCK` never reaches `mock-prod`. See the +[local run instructions](../README.md#run-it-locally) for the exact commands. + +A first attempt at a real `agent-demo/load_driver.py` run against the +hosted tester cluster (after the table above was captured, in the same +session) hit sustained `503 Service Unavailable` from the extract model +(`urchade/gliner_multi-v2.1`) at every concurrency level tried, including +sequential (concurrency=1) requests -- not a capacity limit specific to +concurrent load. A follow-up direct check showed `sie_encode` alone (no +concurrency at all) taking 458 seconds to return, versus roughly a second +earlier in the same session. This points to a transient problem on +Superlinked's shared tester cluster at that moment, not a regression in the +gate or the SDK wiring: `sie_extract`'s own error handling degraded +correctly (returned `[]` rather than raising), just too slowly for +`agent-demo/harness.py`'s 10-second client timeout under any load at all. + +No further load was placed on the cluster once this pattern was clear, out +of courtesy to shared, sponsored compute in a visibly degraded state. + +## Full-stack load test against the recovered hosted cluster + +The hosted tester cluster came back after the outage above, but not into a +steady "always warm" state -- it scales its per-model capacity down to zero +within roughly a minute of no traffic, then re-provisions on the next +request. `sie_score` and `sie_extract` (the two primitives `/v1/gate` +actually calls per request, via `_extra_sie_signals`; `sie_encode` is not +on this request path) each took 0.1-35s to come back from cold before +settling into sub-second responses. This is a real characteristic of a +shared, scale-to-zero tester allocation, not a gate or SDK defect -- +`sie_sdk`'s own transient-error retry handled it transparently in every +case except when a cold re-provision outlasted `agent-demo/harness.py`'s +10-second client timeout. + +**Setup:** `dusk-gate` run locally (not in Docker) with `sie-sdk` installed +temporarily so live SIE calls are actually made (the project's own venv +does not ship `sie-sdk` by default -- it lives in the `sie` extras group, +uninstalled again after this run to keep the venv matching CI); baseline +from `sample-data/baseline.json`; `mock-prod` run locally; full +round trip via `agent-demo/load_driver.py` (`harness.run_scenario` -> +`/v1/gate` -> `mock-prod` on `ALLOW`), 20 requests per concurrency level, +20% poisoned / 80% clean mix, single trial. + +| Concurrency | p50 | p95 | Errors | Verdicts | +|---|---|---|---|---| +| 1 | 294ms | 10008ms | 2/20 | 13 ALLOW, 5 WOULD-BLOCK | +| 3 | 307ms | 474ms | 0/20 | 13 ALLOW, 7 WOULD-BLOCK | +| 5 | 295ms | 317ms | 0/20 | 13 ALLOW, 7 WOULD-BLOCK | + +Correctness held throughout: every `ALLOW` reached `mock-prod` (confirmed +via its `/log`, 46 applied actions across this run and earlier manual +checks) and every poisoned action was `WOULD-BLOCK` in watch mode, never +applied. + +**Reading the errors:** the 2 timeouts at concurrency=1 are cold-provision +blips (a model scaling back to zero between the sparse, sequential +requests at this concurrency, then not re-provisioning inside the 10s +client timeout) -- not a concurrency effect, since concurrency 3 and 5 (more +total request pressure, keeping the cluster continuously warm) both ran +error-free. p50 (294-307ms) is steady and consistent with the earlier +gate-only preliminary probe's p50 (600-670ms; lower here since this run +landed after the extract/score models were already warm going in). + +**Caveats:** n=20 per level, one trial -- enough to confirm the shape (flat +p50 across concurrency, errors tied to idle-driven cold starts rather than +load) but not a high-confidence p95 at concurrency=1. Deliberately kept +small (60 requests total across the sweep) out of courtesy to shared, +sponsored compute. If Superlinked's production SIE tier doesn't scale to +zero this aggressively, the concurrency=1 tail disappears entirely; this is +a property of the tester allocation, worth noting to Superlinked directly +rather than treating as a DUSK-side latency number. diff --git a/examples/agent-action-monitor/docs/sie-primitives.md b/examples/agent-action-monitor/docs/sie-primitives.md new file mode 100644 index 000000000..2b3478adf --- /dev/null +++ b/examples/agent-action-monitor/docs/sie-primitives.md @@ -0,0 +1,105 @@ +# SIE primitives in the agent action gate + +A reference for `examples/agent-action-monitor/README.md`, shaped like +`superlinked/sie`'s existing `stripe-link-fraud` example: a model lineup, +where each primitive is actually used in this codebase, and an honest +account of what the deterministic core still does versus what SIE adds. + +## Model lineup + +| Model | Primitive | Role | +|---|---|---| +| `BAAI/bge-m3` | encode | Embeds an action's description for similarity search against past decisions. | +| `BAAI/bge-reranker-v2-m3` | score | Cross-encoder rerank of the top candidate matches, and of an agent's own history against a new action. | +| `urchade/gliner_multi-v2.1` | extract | Zero-shot extraction of privileged terms (role, privilege, resource, segment, port) from an action's target and change payload, with no training data. | + +All three are verified against the live Superlinked model catalog +(`superlinked.com/models`), not assumed from a family name. + +## Where each primitive is wired in + +- **encode** -- `src/dusk/trace/vector.py`'s `sie_encode()` (wrapped by + `embed_text()`), called live by `/v1/gate` (`src/dusk/api.py`): once per + request to embed the incoming action, and once per verdict to record it + for future lookups. `find_similar_cached()` compares the fresh query + embedding against a bounded, pre-embedded history (capped at 200 entries) + to populate the response's `similar_decision_ids`, without re-embedding + that history on every call. +- **score** -- `src/dusk/trace/vector.py`'s `sie_score()`, used two ways: + reranking the encode-shortlisted candidates for `similar_decision_ids`, + and inside `src/dusk/actions/analyse.py`'s `_semantic_novelty()` to check + a new action's rerank similarity against the acting agent's own raw + baseline history. Its raw cross-encoder output is a logit with no fixed + scale, so `sie_score()` bounds it into `[0, 1]` via sigmoid before a + fixed threshold compares against it -- that bounding is monotonic, not a + calibrated probability (see Known limits below). +- **extract** -- `src/dusk/trace/vector.py`'s `sie_extract()`, used inside + `src/dusk/actions/analyse.py`'s `_extracted_sensitive_terms()` to flag + privileged terms the static frozenset (`_SENSITIVE_TOKENS`/ + `_SENSITIVE_VALUES`) does not already cover. Each extraction keeps its + GLiNER confidence score; terms below `_EXTRACT_CONFIDENCE_FLOOR` (0.5) + are dropped rather than counted, so a low-confidence zero-shot guess + doesn't carry the same weight as one the model was actually sure about. + +## What happens without SIE + +Every one of the three call sites above degrades to a no-op or a +deterministic fallback rather than failing: `sie_encode` falls back to a +hash-based n-gram embedding, `sie_score` and `sie_extract` return `None`/`[]`, +and every downstream signal that depends on them is additive-only, so the +gate's rule-based score is never reduced by their absence. `dusk gate` and +`/v1/gate` work identically without any SIE container running. + +This degrades quickly, not just gracefully: all three calls pass +`wait_for_capacity=False` and a short `provision_timeout_s`, so a model that +isn't warm yet fails in ~1.5s rather than blocking the request while the +SDK's own retry loop waits for it. See the +[local run instructions](../README.md#run-it-locally) for the expected cold +start behavior. + +## Validated against a real SIE cluster + +`tests/test_sie_live_benchmark.py` skips until `DUSK_SIE_ENDPOINT` and, for +authenticated deployments, `SIE_API_KEY` point at a reachable cluster. Run +against Superlinked's hosted tester endpoint, both checks pass: + +- `sie_encode` returns a real 1024-dimension dense vector from `BAAI/bge-m3` + (confirming the model actually loaded and served, not just that the + endpoint answered). +- Precision and recall on the labelled fixture stay at 1.0/1.0, matching the + deterministic-only baseline exactly -- no regression from enabling SIE. +- At least one attack's `reasons` carries a real `SIE rerank` or + `SIE extract` marker, confirming the primitives are actually contributing + a signal over the network, not a no-op that happens to still pass. +- The full test suite passes unchanged with live SIE enabled, confirming + nothing depends on the deterministic fallback path being taken. The exact + test count is intentionally omitted because it changes as coverage grows. + +This is evidence that SIE is load-bearing here ("removing SIE degrades the +result"), not just a claim. + +## Known limits + +- The baseline/attack fixtures used in the benchmark (`lab/actions/ + generate_actions.py`) are synthetic, not real production traffic. +- The deterministic feature checks in `actions/baseline.py` and + `actions/analyse.py` still do the primary anomaly scoring; SIE's three + primitives are an enrichment layer on top, not a replacement for it. This + matches the project's own stance that the core detection logic is not + dependent on any AI model at runtime. +- `sie_score`'s rerank pass only reorders a small shortlist (`top_k`, + default 3) of candidates already retrieved by cosine similarity -- it does + not rerank the full decision history. +- `_SEMANTIC_SIMILARITY_FLOOR` (0.3) is a heuristic cutoff on the + sigmoid-bounded rerank score, not a value derived from an empirical + calibration set. Sigmoid makes the score bounded and monotonic; it does + not make it a calibrated probability that 0.3 has a principled meaning + against. +- The live decision history behind `similar_decision_ids` is in-memory and + capped at 200 entries per gate process -- a demo-scale audit trail, not a + durable store. It resets on restart and is not shared across replicas. +- `sie_extract`'s privileged-term detection is zero-shot: it has not been + evaluated against an adversarial corpus designed to evade GLiNER + specifically, only against the same synthetic fixtures used elsewhere. + The 0.5 confidence floor is a reasonable default, not an empirically + tuned threshold. diff --git a/examples/agent-action-monitor/lab/actions/generate_actions.py b/examples/agent-action-monitor/lab/actions/generate_actions.py new file mode 100644 index 000000000..fb5746651 --- /dev/null +++ b/examples/agent-action-monitor/lab/actions/generate_actions.py @@ -0,0 +1,125 @@ +"""Generate synthetic agent action fixtures (generic JSON format). + +Writes two JSON files, each a list of raw generic-format records (the shape +the generic adapter accepts): + +- ``tests/fixtures/actions_normal.json``: routine, in-scope control-plane + work from a few agents. +- ``tests/fixtures/actions_mixed.json``: the same routine work plus a few + clearly out-of-pattern actions (an agent that only ever edits the corporate + segment suddenly opening a guest-to-restricted firewall rule). + +These exercise ingestion only. This layer assigns no severity; later layers +decide whether the out-of-pattern actions matter. + +Run directly to (re)generate both fixtures:: + + python lab/actions/generate_actions.py +""" + +from __future__ import annotations + +import json +import os +import sys +from datetime import UTC, datetime, timedelta +from typing import Any + +BASE_TIME = datetime(2023, 11, 14, 22, 13, 20, tzinfo=UTC) +FIXTURE_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "tests", "fixtures") + + +def _action( + offset: int, + agent_id: str, + action_type: str, + target: str, + before: Any = None, + after: Any = None, +) -> dict[str, Any]: + """Build one generic-format record with an evenly spaced timestamp.""" + return { + "agent_id": agent_id, + "timestamp": (BASE_TIME + timedelta(minutes=offset)).isoformat(), + "action_type": action_type, + "target": target, + "change": {"before": before, "after": after}, + "source": "generic", + "raw_ref": f"evt-{offset:04d}", + } + + +def normal_actions() -> list[dict[str, Any]]: + """Routine, in-scope actions from three agents.""" + return [ + _action(0, "netops-agent", "firewall_rule_change", "fw-corp-https", after={"port": 443}), + _action(1, "netops-agent", "firewall_rule_change", "fw-corp-dns", after={"port": 53}), + _action( + 2, "netops-agent", "route_change", "rt-corp-to-dmz", after={"next_hop": "10.0.0.1"} + ), + _action(3, "netops-agent", "route_change", "rt-corp-default"), + _action(4, "netops-agent", "firewall_rule_change", "fw-temp-debug", before={"port": 22}), + _action( + 5, "segment-agent", "segment_change", "seg-corporate", after={"cidr": "10.0.10.0/24"} + ), + _action( + 6, "segment-agent", "segment_change", "seg-corporate", after={"cidr": "10.0.10.0/23"} + ), + _action(7, "segment-agent", "segment_change", "seg-corp-staging"), + _action(8, "iam-agent", "role_assignment", "ra-netops-reader", after={"role": "reader"}), + _action(9, "iam-agent", "role_assignment", "ra-segment-contributor"), + _action(10, "netops-agent", "firewall_rule_change", "fw-corp-https", after={"port": 443}), + _action(11, "netops-agent", "route_change", "rt-corp-to-monitoring"), + _action(12, "segment-agent", "segment_change", "seg-corporate"), + _action(13, "iam-agent", "role_assignment", "ra-stale-temp", before={"role": "reader"}), + _action(14, "netops-agent", "port_change", "fw-corp-dns", after={"port": 53}), + ] + + +def out_of_pattern_actions() -> list[dict[str, Any]]: + """Actions that break each agent's established pattern.""" + return [ + # The segment agent has only ever touched segments; now it opens a + # firewall path from guest into the restricted segment. + _action( + 15, + "segment-agent", + "firewall_rule_change", + "fw-guest-to-restricted", + after={"src": "10.0.40.0/24", "dst": "10.0.99.0/24", "port": 445}, + ), + # The IAM agent grants itself a broad owner role. + _action(16, "iam-agent", "role_assignment", "ra-iam-owner-self", after={"role": "owner"}), + # The netops agent reassigns a restricted segment it never manages. + _action(17, "netops-agent", "segment_change", "seg-restricted"), + ] + + +def generate(directory: str = FIXTURE_DIR) -> tuple[str, str]: + """Write both fixtures into ``directory`` and return their paths.""" + os.makedirs(os.path.abspath(directory), exist_ok=True) + normal = normal_actions() + mixed = normal + out_of_pattern_actions() + + normal_path = os.path.join(directory, "actions_normal.json") + mixed_path = os.path.join(directory, "actions_mixed.json") + with open(normal_path, "w", encoding="utf-8") as handle: + json.dump(normal, handle, indent=2) + with open(mixed_path, "w", encoding="utf-8") as handle: + json.dump(mixed, handle, indent=2) + return os.path.abspath(normal_path), os.path.abspath(mixed_path) + + +def main() -> int: + try: + normal_path, mixed_path = generate() + except OSError as exc: + print(f"Failed to write fixtures: {exc}", file=sys.stderr) + return 1 + print(f"Wrote {normal_path}") + print(f"Wrote {mixed_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/agent-action-monitor/mock-prod/Dockerfile b/examples/agent-action-monitor/mock-prod/Dockerfile new file mode 100644 index 000000000..296c93fcc --- /dev/null +++ b/examples/agent-action-monitor/mock-prod/Dockerfile @@ -0,0 +1,23 @@ +# mock-prod: dummy downstream controller/DB target for the DUSK example. +FROM python:3.11-slim + +WORKDIR /app + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt \ + && useradd --create-home --uid 1000 mockprod \ + && chown -R mockprod:mockprod /app + +USER mockprod + +COPY app.py . + +ENV MOCK_PROD_HOST=0.0.0.0 \ + MOCK_PROD_PORT=9000 + +EXPOSE 9000 + +HEALTHCHECK --interval=10s --timeout=3s --start-period=5s --retries=3 \ + CMD ["python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:9000/health', timeout=3)"] + +CMD ["python", "app.py"] diff --git a/examples/agent-action-monitor/mock-prod/app.py b/examples/agent-action-monitor/mock-prod/app.py new file mode 100644 index 000000000..8d27bfd26 --- /dev/null +++ b/examples/agent-action-monitor/mock-prod/app.py @@ -0,0 +1,55 @@ +"""In-memory downstream target used to verify gate enforcement.""" + +from __future__ import annotations + +import logging +import os +from datetime import UTC, datetime +from typing import Any + +from flask import Flask, request + +logging.basicConfig(level=logging.INFO, format="%(asctime)s mock-prod: %(message)s") +logger = logging.getLogger("mock-prod") + +app = Flask(__name__) + +#: In-memory log of applied actions, most recent last. Demo-scale only. +applied_log: list[dict[str, Any]] = [] + + +@app.get("/health") +def health() -> tuple[dict[str, str], int]: + return {"status": "ok"}, 200 + + +@app.post("/apply") +def apply() -> tuple[dict[str, Any], int]: + action = request.get_json(force=True, silent=True) + if not isinstance(action, dict): + return {"error": "expected a JSON object"}, 400 + + record = { + "received_at": datetime.now(UTC).isoformat(), + "agent_id": action.get("agent_id"), + "action_type": action.get("action_type"), + "target": action.get("target"), + } + applied_log.append(record) + logger.info("applied: %s", record) + return {"status": "applied", **record}, 200 + + +@app.get("/log") +def log() -> tuple[dict[str, Any], int]: + return {"count": len(applied_log), "entries": applied_log}, 200 + + +def run() -> None: + port = int(os.getenv("MOCK_PROD_PORT", "9000")) + host = os.getenv("MOCK_PROD_HOST", "127.0.0.1") + app.run(host=host, port=port) + + +if __name__ == "__main__": + run() diff --git a/examples/agent-action-monitor/mock-prod/requirements.txt b/examples/agent-action-monitor/mock-prod/requirements.txt new file mode 100644 index 000000000..001e7c4af --- /dev/null +++ b/examples/agent-action-monitor/mock-prod/requirements.txt @@ -0,0 +1 @@ +flask>=3.0 diff --git a/examples/agent-action-monitor/mock-prod/test_app.py b/examples/agent-action-monitor/mock-prod/test_app.py new file mode 100644 index 000000000..603176a94 --- /dev/null +++ b/examples/agent-action-monitor/mock-prod/test_app.py @@ -0,0 +1,43 @@ +"""Tests for the mock-prod dummy downstream target.""" + +from __future__ import annotations + +import pytest +from app import app, applied_log + + +@pytest.fixture(autouse=True) +def _clear_log(): + applied_log.clear() + yield + applied_log.clear() + + +@pytest.fixture +def client(): + app.testing = True + return app.test_client() + + +def test_health(client): + resp = client.get("/health") + assert resp.status_code == 200 + + +def test_apply_logs_the_action(client): + action = {"agent_id": "agent-1", "action_type": "route_change", "target": "rt-123"} + resp = client.post("/apply", json=action) + assert resp.status_code == 200 + body = resp.get_json() + assert body["status"] == "applied" + assert body["agent_id"] == "agent-1" + + log_resp = client.get("/log") + log_body = log_resp.get_json() + assert log_body["count"] == 1 + assert log_body["entries"][0]["target"] == "rt-123" + + +def test_apply_rejects_non_object_body(client): + resp = client.post("/apply", json=["not", "an", "object"]) + assert resp.status_code == 400 diff --git a/examples/agent-action-monitor/n8n/Dockerfile b/examples/agent-action-monitor/n8n/Dockerfile new file mode 100644 index 000000000..ee6a04084 --- /dev/null +++ b/examples/agent-action-monitor/n8n/Dockerfile @@ -0,0 +1,16 @@ +# n8n pre-loaded with the three DUSK gate webhooks (decision/report/alert), +# active from container start -- no manual UI import step, no external +# service in the workflow itself (each webhook just responds immediately). +FROM n8nio/n8n:2.29.10 + +COPY dusk-webhooks.json /dusk-webhooks.json +COPY docker-entrypoint.sh /docker-entrypoint-dusk.sh + +USER root +RUN chmod +x /docker-entrypoint-dusk.sh && chown node:node /dusk-webhooks.json +USER node + +HEALTHCHECK --interval=10s --timeout=3s --start-period=20s --retries=5 \ + CMD ["wget", "-qO-", "http://localhost:5678/healthz"] + +ENTRYPOINT ["tini", "--", "/docker-entrypoint-dusk.sh"] diff --git a/examples/agent-action-monitor/n8n/docker-entrypoint.sh b/examples/agent-action-monitor/n8n/docker-entrypoint.sh new file mode 100644 index 000000000..6c416c973 --- /dev/null +++ b/examples/agent-action-monitor/n8n/docker-entrypoint.sh @@ -0,0 +1,7 @@ +#!/bin/sh +set -e + +n8n import:workflow --input=/dusk-webhooks.json +n8n publish:workflow --id=dusk-gate-webhooks + +exec n8n start diff --git a/examples/agent-action-monitor/n8n/dusk-webhooks.json b/examples/agent-action-monitor/n8n/dusk-webhooks.json new file mode 100644 index 000000000..db0aa099b --- /dev/null +++ b/examples/agent-action-monitor/n8n/dusk-webhooks.json @@ -0,0 +1,59 @@ +[ + { + "id": "dusk-gate-webhooks", + "name": "DUSK gate webhooks", + "active": true, + "nodes": [ + { + "parameters": { + "httpMethod": "POST", + "path": "dusk-decision", + "responseMode": "onReceived", + "responseData": "{\"status\":\"received\"}", + "options": {} + }, + "id": "n1", + "name": "decision", + "type": "n8n-nodes-base.webhook", + "typeVersion": 2, + "position": [400, 200], + "webhookId": "dusk-decision" + }, + { + "parameters": { + "httpMethod": "POST", + "path": "dusk-report", + "responseMode": "onReceived", + "responseData": "{\"status\":\"received\"}", + "options": {} + }, + "id": "n2", + "name": "report", + "type": "n8n-nodes-base.webhook", + "typeVersion": 2, + "position": [400, 400], + "webhookId": "dusk-report" + }, + { + "parameters": { + "httpMethod": "POST", + "path": "dusk-alert", + "responseMode": "onReceived", + "responseData": "{\"status\":\"received\"}", + "options": {} + }, + "id": "n3", + "name": "alert", + "type": "n8n-nodes-base.webhook", + "typeVersion": 2, + "position": [400, 600], + "webhookId": "dusk-alert" + } + ], + "connections": {}, + "settings": {}, + "staticData": null, + "pinData": {}, + "tags": [] + } +] diff --git a/examples/agent-action-monitor/pyproject.toml b/examples/agent-action-monitor/pyproject.toml new file mode 100644 index 000000000..54d4026b8 --- /dev/null +++ b/examples/agent-action-monitor/pyproject.toml @@ -0,0 +1,80 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "dusk-agent-action-monitor" +version = "0.1.0" +description = "DUSK's agent action gate, judging a proposed agent action against a per-agent behavioural baseline, with Superlinked SIE surfacing the anomalies." +readme = "README.md" +requires-python = ">=3.11" +license = "Apache-2.0" +authors = [{ name = "Ritik Sah" }, { name = "Tanvir Farhad" }] +keywords = ["security", "ai-agents", "behavioral-analysis", "superlinked"] + +dependencies = [ + "PyYAML>=6.0", + "flask>=3.0", + "flask-cors>=4.0", + "python-dotenv>=1.0", +] + +[project.optional-dependencies] +sie = [ + # Tested compatibility pair: SDK 0.6.17 with server v0.4.1-cpu-default. + "sie-sdk==0.6.17", +] +dev = [ + "pytest>=7.4.0", + "pytest-cov>=4.1.0", + "pytest-env>=1.1.0", + "ruff>=0.4.4", + "mypy>=1.10.0", + "bandit[toml]>=1.7.8", + "types-PyYAML>=6.0", + "requests>=2.31", # agent-demo/harness.py's HTTP calls to the gate and mock-prod +] + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.pytest.ini_options] +testpaths = ["tests", "agent-demo", "mock-prod"] +env = ["PYTHONIOENCODING=utf-8"] + +[tool.ruff] +target-version = "py311" +line-length = 100 + +[tool.ruff.lint] +select = ["E", "W", "F", "I", "N", "S", "B", "C90", "UP", "ANN"] + +[tool.ruff.lint.per-file-ignores] +"tests/*" = ["S", "ANN"] +"lab/*" = ["S", "ANN"] +"agent-demo/test_*.py" = ["S", "ANN"] +"mock-prod/test_*.py" = ["S", "ANN"] +# Bare-name statements are vulture's own documented whitelist syntax, not a +# real code smell -- see scripts/vulture_whitelist.py's module docstring. +"scripts/vulture_whitelist.py" = ["B018"] + +[tool.mypy] +python_version = "3.11" +strict = true +warn_return_any = true +warn_unused_configs = true +disallow_untyped_defs = true +disallow_any_generics = true + +[[tool.mypy.overrides]] +module = ["dusk.api"] +ignore_missing_imports = true +disable_error_code = ["type-var", "untyped-decorator", "unused-ignore", "import-untyped", "import-not-found"] + +[[tool.mypy.overrides]] +module = ["dusk.trace.*"] +ignore_missing_imports = true +disable_error_code = ["import-untyped", "attr-defined", "unused-ignore"] + +[tool.bandit] +exclude_dirs = ["tests"] diff --git a/examples/agent-action-monitor/sample-data/baseline.json b/examples/agent-action-monitor/sample-data/baseline.json new file mode 100644 index 000000000..975fa8f4a --- /dev/null +++ b/examples/agent-action-monitor/sample-data/baseline.json @@ -0,0 +1,202 @@ +[ + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:13:20+00:00", + "action_type": "firewall_rule_change", + "target": "fw-corp-https", + "change": { + "before": null, + "after": { + "port": 443 + } + }, + "source": "generic", + "raw_ref": "evt-0000" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:14:20+00:00", + "action_type": "firewall_rule_change", + "target": "fw-corp-dns", + "change": { + "before": null, + "after": { + "port": 53 + } + }, + "source": "generic", + "raw_ref": "evt-0001" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:15:20+00:00", + "action_type": "route_change", + "target": "rt-corp-to-dmz", + "change": { + "before": null, + "after": { + "next_hop": "10.0.0.1" + } + }, + "source": "generic", + "raw_ref": "evt-0002" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:16:20+00:00", + "action_type": "route_change", + "target": "rt-corp-default", + "change": { + "before": null, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0003" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:17:20+00:00", + "action_type": "firewall_rule_change", + "target": "fw-temp-debug", + "change": { + "before": { + "port": 22 + }, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0004" + }, + { + "agent_id": "segment-agent", + "timestamp": "2023-11-14T22:18:20+00:00", + "action_type": "segment_change", + "target": "seg-corporate", + "change": { + "before": null, + "after": { + "cidr": "10.0.10.0/24" + } + }, + "source": "generic", + "raw_ref": "evt-0005" + }, + { + "agent_id": "segment-agent", + "timestamp": "2023-11-14T22:19:20+00:00", + "action_type": "segment_change", + "target": "seg-corporate", + "change": { + "before": null, + "after": { + "cidr": "10.0.10.0/23" + } + }, + "source": "generic", + "raw_ref": "evt-0006" + }, + { + "agent_id": "segment-agent", + "timestamp": "2023-11-14T22:20:20+00:00", + "action_type": "segment_change", + "target": "seg-corp-staging", + "change": { + "before": null, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0007" + }, + { + "agent_id": "iam-agent", + "timestamp": "2023-11-14T22:21:20+00:00", + "action_type": "role_assignment", + "target": "ra-netops-reader", + "change": { + "before": null, + "after": { + "role": "reader" + } + }, + "source": "generic", + "raw_ref": "evt-0008" + }, + { + "agent_id": "iam-agent", + "timestamp": "2023-11-14T22:22:20+00:00", + "action_type": "role_assignment", + "target": "ra-segment-contributor", + "change": { + "before": null, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0009" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:23:20+00:00", + "action_type": "firewall_rule_change", + "target": "fw-corp-https", + "change": { + "before": null, + "after": { + "port": 443 + } + }, + "source": "generic", + "raw_ref": "evt-0010" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:24:20+00:00", + "action_type": "route_change", + "target": "rt-corp-to-monitoring", + "change": { + "before": null, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0011" + }, + { + "agent_id": "segment-agent", + "timestamp": "2023-11-14T22:25:20+00:00", + "action_type": "segment_change", + "target": "seg-corporate", + "change": { + "before": null, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0012" + }, + { + "agent_id": "iam-agent", + "timestamp": "2023-11-14T22:26:20+00:00", + "action_type": "role_assignment", + "target": "ra-stale-temp", + "change": { + "before": { + "role": "reader" + }, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0013" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:27:20+00:00", + "action_type": "port_change", + "target": "fw-corp-dns", + "change": { + "before": null, + "after": { + "port": 53 + } + }, + "source": "generic", + "raw_ref": "evt-0014" + } +] \ No newline at end of file diff --git a/examples/agent-action-monitor/sample-data/check-mixed.json b/examples/agent-action-monitor/sample-data/check-mixed.json new file mode 100644 index 000000000..a025927c8 --- /dev/null +++ b/examples/agent-action-monitor/sample-data/check-mixed.json @@ -0,0 +1,244 @@ +[ + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:13:20+00:00", + "action_type": "firewall_rule_change", + "target": "fw-corp-https", + "change": { + "before": null, + "after": { + "port": 443 + } + }, + "source": "generic", + "raw_ref": "evt-0000" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:14:20+00:00", + "action_type": "firewall_rule_change", + "target": "fw-corp-dns", + "change": { + "before": null, + "after": { + "port": 53 + } + }, + "source": "generic", + "raw_ref": "evt-0001" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:15:20+00:00", + "action_type": "route_change", + "target": "rt-corp-to-dmz", + "change": { + "before": null, + "after": { + "next_hop": "10.0.0.1" + } + }, + "source": "generic", + "raw_ref": "evt-0002" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:16:20+00:00", + "action_type": "route_change", + "target": "rt-corp-default", + "change": { + "before": null, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0003" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:17:20+00:00", + "action_type": "firewall_rule_change", + "target": "fw-temp-debug", + "change": { + "before": { + "port": 22 + }, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0004" + }, + { + "agent_id": "segment-agent", + "timestamp": "2023-11-14T22:18:20+00:00", + "action_type": "segment_change", + "target": "seg-corporate", + "change": { + "before": null, + "after": { + "cidr": "10.0.10.0/24" + } + }, + "source": "generic", + "raw_ref": "evt-0005" + }, + { + "agent_id": "segment-agent", + "timestamp": "2023-11-14T22:19:20+00:00", + "action_type": "segment_change", + "target": "seg-corporate", + "change": { + "before": null, + "after": { + "cidr": "10.0.10.0/23" + } + }, + "source": "generic", + "raw_ref": "evt-0006" + }, + { + "agent_id": "segment-agent", + "timestamp": "2023-11-14T22:20:20+00:00", + "action_type": "segment_change", + "target": "seg-corp-staging", + "change": { + "before": null, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0007" + }, + { + "agent_id": "iam-agent", + "timestamp": "2023-11-14T22:21:20+00:00", + "action_type": "role_assignment", + "target": "ra-netops-reader", + "change": { + "before": null, + "after": { + "role": "reader" + } + }, + "source": "generic", + "raw_ref": "evt-0008" + }, + { + "agent_id": "iam-agent", + "timestamp": "2023-11-14T22:22:20+00:00", + "action_type": "role_assignment", + "target": "ra-segment-contributor", + "change": { + "before": null, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0009" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:23:20+00:00", + "action_type": "firewall_rule_change", + "target": "fw-corp-https", + "change": { + "before": null, + "after": { + "port": 443 + } + }, + "source": "generic", + "raw_ref": "evt-0010" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:24:20+00:00", + "action_type": "route_change", + "target": "rt-corp-to-monitoring", + "change": { + "before": null, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0011" + }, + { + "agent_id": "segment-agent", + "timestamp": "2023-11-14T22:25:20+00:00", + "action_type": "segment_change", + "target": "seg-corporate", + "change": { + "before": null, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0012" + }, + { + "agent_id": "iam-agent", + "timestamp": "2023-11-14T22:26:20+00:00", + "action_type": "role_assignment", + "target": "ra-stale-temp", + "change": { + "before": { + "role": "reader" + }, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0013" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:27:20+00:00", + "action_type": "port_change", + "target": "fw-corp-dns", + "change": { + "before": null, + "after": { + "port": 53 + } + }, + "source": "generic", + "raw_ref": "evt-0014" + }, + { + "agent_id": "segment-agent", + "timestamp": "2023-11-14T22:28:20+00:00", + "action_type": "firewall_rule_change", + "target": "fw-guest-to-restricted", + "change": { + "before": null, + "after": { + "src": "10.0.40.0/24", + "dst": "10.0.99.0/24", + "port": 445 + } + }, + "source": "generic", + "raw_ref": "evt-0015" + }, + { + "agent_id": "iam-agent", + "timestamp": "2023-11-14T22:29:20+00:00", + "action_type": "role_assignment", + "target": "ra-iam-owner-self", + "change": { + "before": null, + "after": { + "role": "owner" + } + }, + "source": "generic", + "raw_ref": "evt-0016" + }, + { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:30:20+00:00", + "action_type": "segment_change", + "target": "seg-restricted", + "change": { + "before": null, + "after": null + }, + "source": "generic", + "raw_ref": "evt-0017" + } +] \ No newline at end of file diff --git a/examples/agent-action-monitor/scripts/vulture_whitelist.py b/examples/agent-action-monitor/scripts/vulture_whitelist.py new file mode 100644 index 000000000..25c91ec8b --- /dev/null +++ b/examples/agent-action-monitor/scripts/vulture_whitelist.py @@ -0,0 +1,29 @@ +"""Known-legitimate "unused" code, per vulture's whitelist convention. + +Referenced by name only, never imported or executed: this file exists to be +scanned by vulture (`vulture src/ tests/ agent-demo/ mock-prod/ +scripts/vulture_whitelist.py`), not to run. Every entry here was checked +against a real cross-reference search before being added. + +- ActionGate.evaluate_all(): batch evaluation, used by the root repo's + offline `dusk gate` CLI (a separate package -- this example's live + /v1/gate handler processes one action per request via .evaluate(), so it + never needs the batch method). Kept for parity with the root copy's + public API, not dead code that was meant to be removed. +- config.py's set_config(): a real public setter with no caller in this + package yet (tests use monkeypatch/env vars instead), kept for parity + with get_config()/reset_config() as the documented override path. +- vector.py's SimilarDecision.similarity: set at construction (the + cosine/rerank score each match was found at) but never read back by + api.py, which only forwards each match's .id in similar_decision_ids. + Kept on the dataclass as diagnostic data for a caller that wants the + actual similarity score, not just which decisions matched. +""" + +from dusk.actions.verdict import ActionGate +from dusk.config import set_config +from dusk.trace.vector import SimilarDecision + +ActionGate.evaluate_all +set_config +SimilarDecision.similarity diff --git a/examples/agent-action-monitor/src/dusk/__init__.py b/examples/agent-action-monitor/src/dusk/__init__.py new file mode 100644 index 000000000..b4573a452 --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/__init__.py @@ -0,0 +1,10 @@ +"""Dusk, behavioral threat detection for agentic networks. + +Dusk watches how AI agents move through a network and flags the +machine-paced, systematic patterns that signal a hijacked or poisoned +agent acting against the network it controls. +""" + +__version__ = "0.1.0" + +__all__ = ["__version__"] diff --git a/examples/agent-action-monitor/src/dusk/actions/__init__.py b/examples/agent-action-monitor/src/dusk/actions/__init__.py new file mode 100644 index 000000000..e66b6ce65 --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/__init__.py @@ -0,0 +1,36 @@ +"""The agent action layer: ingest, baseline, analyse, and gate agent actions. + +Agent actions arrive from whatever controller the environment uses (a cloud +network API, an SDN controller, a policy endpoint) and are normalised into one +canonical :class:`~dusk.actions.event.AgentAction` event. The gate then learns +each agent's normal behaviour, scores a new action against that baseline, and +renders a verdict (ALLOW, WOULD-BLOCK, BLOCK) with reasons and MITRE mappings. +Source-specific shapes are handled by adapters (see :mod:`dusk.actions.adapters`). +""" + +from dusk.actions.adapters.base import AdapterError, SourceAdapter +from dusk.actions.analyse import AnalysisResult, analyse +from dusk.actions.baseline import AgentProfile, Baseline +from dusk.actions.event import KNOWN_ACTION_TYPES, AgentAction +from dusk.actions.ingest import ingest_file +from dusk.actions.normaliser import known_sources, normalise_record +from dusk.actions.verdict import ALLOW, BLOCK, WOULD_BLOCK, ActionGate, GateVerdict + +__all__ = [ + "ALLOW", + "BLOCK", + "KNOWN_ACTION_TYPES", + "WOULD_BLOCK", + "ActionGate", + "AdapterError", + "AgentAction", + "AgentProfile", + "AnalysisResult", + "Baseline", + "GateVerdict", + "SourceAdapter", + "analyse", + "ingest_file", + "known_sources", + "normalise_record", +] diff --git a/examples/agent-action-monitor/src/dusk/actions/adapters/__init__.py b/examples/agent-action-monitor/src/dusk/actions/adapters/__init__.py new file mode 100644 index 000000000..e95a05c54 --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/adapters/__init__.py @@ -0,0 +1,5 @@ +"""Adapters from vendor-specific records to AgentAction.""" + +from dusk.actions.adapters.base import AdapterError, SourceAdapter + +__all__ = ["AdapterError", "SourceAdapter"] diff --git a/examples/agent-action-monitor/src/dusk/actions/adapters/azure.py b/examples/agent-action-monitor/src/dusk/actions/adapters/azure.py new file mode 100644 index 000000000..fc3b67aa9 --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/adapters/azure.py @@ -0,0 +1,119 @@ +"""Azure activity-log adapter. + +Maps an Azure Monitor activity-log record into the canonical AgentAction. + +Field assumptions (documented because Azure's activity-log shape varies by +API version, and these are the fields this adapter relies on): + +- ``operationName``: the operation. In the activity-log schema this can be a + plain string or an object ``{"value": ..., "localizedValue": ...}``. This + adapter accepts both and falls back to ``authorization.action`` (the RBAC + action, for example ``"Microsoft.Network/networkSecurityGroups/write"``). +- ``caller``: the acting identity (a user principal name or object id). Used + as ``agent_id``. +- ``resourceId``: the affected resource. Used as ``target``. +- ``eventTimestamp``: ISO 8601 time of the event. A trailing ``Z`` is + accepted. Used as ``timestamp``. +- ``properties``: optional, may carry ``before``/``after`` deltas. Used to + build ``change``. +- ``correlationId`` / ``eventDataId``: opaque reference to the record. Used as + ``raw_ref``. + +Operation-name to action_type mapping (substring, case-insensitive): + +- networkSecurityGroups or securityRules -> ``firewall_rule_change`` +- routeTables or routes -> ``route_change`` +- virtualNetworks or subnets -> ``segment_change`` +- roleAssignments -> ``role_assignment`` +- anything else -> ``unknown`` + +This adapter normalises shape only. It assigns no severity. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from dusk.actions.adapters.base import AdapterError, SourceAdapter +from dusk.actions.event import AgentAction + +#: Substrings of the operation name mapped to a canonical action_type. +_ACTION_TYPE_RULES: tuple[tuple[tuple[str, ...], str], ...] = ( + (("networksecuritygroups", "securityrules"), "firewall_rule_change"), + (("routetables", "routes"), "route_change"), + (("virtualnetworks", "subnets"), "segment_change"), + (("roleassignments",), "role_assignment"), +) + + +def _operation_name(raw: dict[str, Any]) -> str | None: + """Extract the operation name as a string, accepting both Azure shapes.""" + op = raw.get("operationName") + if isinstance(op, dict): + op = op.get("value") or op.get("localizedValue") + if not op: + authorization = raw.get("authorization") + if isinstance(authorization, dict): + op = authorization.get("action") + return op if isinstance(op, str) and op else None + + +def _action_type(operation: str | None) -> str: + """Classify the canonical action_type from the operation name.""" + if not operation: + return "unknown" + lowered = operation.lower() + for needles, action_type in _ACTION_TYPE_RULES: + if any(needle in lowered for needle in needles): + return action_type + return "unknown" + + +class AzureAdapter(SourceAdapter): + """Adapter for Azure Monitor activity-log records.""" + + source = "azure" + + def parse(self, raw: dict[str, Any]) -> AgentAction: + """Map one Azure activity-log record into an :class:`AgentAction`. + + Args: + raw: One Azure activity-log record. + + Returns: + The canonical :class:`AgentAction`. + + Raises: + AdapterError: If the record lacks the identity, target, or + timestamp needed to form a valid AgentAction. + """ + operation = _operation_name(raw) + action_type = _action_type(operation) + + properties = raw.get("properties") + if isinstance(properties, dict): + change: dict[str, Any] = { + "before": properties.get("before"), + "after": properties.get("after"), + } + else: + change = {"before": None, "after": None} + + event_timestamp = raw.get("eventTimestamp") + try: + if not isinstance(event_timestamp, str): + raise AdapterError("Azure record has no string eventTimestamp") + # Azure may emit a trailing 'Z'; normalise it to an explicit offset. + timestamp = datetime.fromisoformat(event_timestamp.replace("Z", "+00:00")) + return AgentAction( + agent_id=raw.get("caller") or "", + timestamp=timestamp, + action_type=action_type, + target=raw.get("resourceId") or "", + change=change, + source=self.source, + raw_ref=raw.get("correlationId") or raw.get("eventDataId"), + ) + except (TypeError, ValueError) as exc: + raise AdapterError(f"Malformed Azure record: {exc}") from exc diff --git a/examples/agent-action-monitor/src/dusk/actions/adapters/base.py b/examples/agent-action-monitor/src/dusk/actions/adapters/base.py new file mode 100644 index 000000000..4674f79f1 --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/adapters/base.py @@ -0,0 +1,49 @@ +"""Base contract for source adapters. + +A :class:`SourceAdapter` maps one source's raw record into the canonical +:class:`~dusk.actions.event.AgentAction`. The contract is deliberate: + +- Map source-specific operation names onto the canonical fields. Where a + field is optional and absent, apply a sensible default (an empty + ``details`` mapping, for example) rather than failing. +- Raise :class:`AdapterError` only when the record is malformed enough that + it cannot yield a valid AgentAction (missing identity, missing target, + unparseable timestamp). One bad record is the caller's to skip; a + structurally impossible record is the adapter's to reject. + +Adapters normalise shape only. They assign no severity, score, or verdict; +those belong to later layers. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from dusk.actions.event import AgentAction + + +class AdapterError(ValueError): + """Raised when a raw source record cannot yield a valid AgentAction.""" + + +class SourceAdapter(ABC): + """Abstract base for a single control-plane source adapter.""" + + #: Canonical source name this adapter handles (for example ``"azure"``). + source: str = "base" + + @abstractmethod + def parse(self, raw: dict[str, Any]) -> AgentAction: + """Map one raw source record into an :class:`AgentAction`. + + Args: + raw: One record in the source's native shape. + + Returns: + The canonical :class:`AgentAction`. + + Raises: + AdapterError: If the record cannot yield a valid AgentAction. + """ + raise NotImplementedError diff --git a/examples/agent-action-monitor/src/dusk/actions/adapters/bedrock.py b/examples/agent-action-monitor/src/dusk/actions/adapters/bedrock.py new file mode 100644 index 000000000..91f17c88a --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/adapters/bedrock.py @@ -0,0 +1,120 @@ +"""Normalize proposed Bedrock Converse tool calls before execution. + +The harness supplies agent identity and timestamp because Bedrock's +``toolUse`` block contains only the tool name, arguments, and call ID. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from dusk.actions.adapters.base import AdapterError, SourceAdapter +from dusk.actions.event import AgentAction + +#: Substrings of the tool name mapped to a canonical action_type. +_ACTION_TYPE_RULES: tuple[tuple[tuple[str, ...], str], ...] = ( + (("firewall", "securitygroup"), "firewall_rule_change"), + (("route",), "route_change"), + (("segment", "subnet", "vpc"), "segment_change"), + (("role", "permission"), "role_assignment"), + (("port",), "port_change"), +) + + +def _action_type(tool_name: str | None) -> str: + """Classify the canonical action_type from the tool name.""" + if not tool_name: + return "unknown" + lowered = tool_name.lower() + for needles, action_type in _ACTION_TYPE_RULES: + if any(needle in lowered for needle in needles): + return action_type + return "unknown" + + +class BedrockAdapter(SourceAdapter): + """Adapter for a proposed Bedrock Converse API tool-call.""" + + source = "bedrock" + + def parse(self, raw: dict[str, Any]) -> AgentAction: + """Map a raw record already carrying agent_id/timestamp into an AgentAction. + + Satisfies the :class:`SourceAdapter` contract for registry lookup + and the generic ingest path. ``raw`` must additionally carry + ``agent_id`` and ``timestamp`` (the harness's job to attach -- + see :func:`parse_tool_use` for the primary entry point used by + agent-demo). + + Args: + raw: A dict with ``tool_use`` (Bedrock's toolUse block) plus + ``agent_id`` and ``timestamp``. + + Returns: + The canonical :class:`AgentAction`. + + Raises: + AdapterError: If the tool-call, identity, or timestamp is + missing or malformed. + """ + tool_use = raw.get("tool_use") + if not isinstance(tool_use, dict): + raise AdapterError("Bedrock record missing 'tool_use' block") + try: + timestamp = raw["timestamp"] + if isinstance(timestamp, str): + timestamp = datetime.fromisoformat(timestamp) + return self.parse_tool_use(tool_use, agent_id=raw["agent_id"], timestamp=timestamp) + except KeyError as exc: + raise AdapterError(f"Bedrock record missing required field: {exc}") from exc + except (TypeError, ValueError) as exc: + raise AdapterError(f"Malformed Bedrock record: {exc}") from exc + + def parse_tool_use( + self, tool_use: dict[str, Any], *, agent_id: str, timestamp: datetime + ) -> AgentAction: + """Map one Bedrock toolUse block into an :class:`AgentAction`. + + Args: + tool_use: The ``toolUse`` content block from a Converse API + response (``name``, ``input``, ``toolUseId``). + agent_id: Identity of the agent that made the call. Supplied + by the harness, not present in the toolUse block itself. + timestamp: When the call was made. Must be timezone-aware. + + Returns: + The canonical :class:`AgentAction`. + + Raises: + AdapterError: If the tool-call lacks the name or target input + needed to form a valid AgentAction. + """ + tool_name = tool_use.get("name") + action_type = _action_type(tool_name if isinstance(tool_name, str) else None) + + tool_input = tool_use.get("input") + if not isinstance(tool_input, dict): + raise AdapterError("Bedrock toolUse block missing 'input'") + + target = tool_input.get("target") + if not target: + raise AdapterError("Bedrock toolUse input missing 'target'") + + change = { + "before": tool_input.get("before"), + "after": tool_input.get("after"), + } + + try: + return AgentAction( + agent_id=agent_id, + timestamp=timestamp, + action_type=action_type, + target=target, + change=change, + source=self.source, + raw_ref=tool_use.get("toolUseId"), + ) + except ValueError as exc: + raise AdapterError(f"Malformed Bedrock tool-call: {exc}") from exc diff --git a/examples/agent-action-monitor/src/dusk/actions/adapters/generic.py b/examples/agent-action-monitor/src/dusk/actions/adapters/generic.py new file mode 100644 index 000000000..92885f02c --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/adapters/generic.py @@ -0,0 +1,60 @@ +"""Generic adapter for vendor-neutral action records. + +The generic record already uses the canonical field names (``agent_id``, +``timestamp``, ``action_type``, ``target``, ``change``, ``source``, optional +``raw_ref``). This is the path for any source we have not written a dedicated +adapter for yet, and the path the synthetic generator uses. The adapter +validates and builds an :class:`~dusk.actions.event.AgentAction`, surfacing +any problem as an :class:`AdapterError`. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from dusk.actions.adapters.base import AdapterError, SourceAdapter +from dusk.actions.event import AgentAction + + +class GenericAdapter(SourceAdapter): + """Adapter for records already shaped like the canonical schema.""" + + source = "generic" + + def parse(self, raw: dict[str, Any]) -> AgentAction: + """Validate and build an :class:`AgentAction` from a generic record. + + Args: + raw: A record whose keys match the canonical AgentAction fields. + ``timestamp`` may be a datetime or an ISO 8601 string. + + Returns: + The canonical :class:`AgentAction`. + + Raises: + AdapterError: If a required field is missing or any value is + invalid (empty identity or target, naive timestamp, unknown + action_type). + """ + try: + raw_timestamp = raw["timestamp"] + timestamp = ( + datetime.fromisoformat(raw_timestamp) + if isinstance(raw_timestamp, str) + else raw_timestamp + ) + change = raw.get("change") or {"before": None, "after": None} + return AgentAction( + agent_id=raw["agent_id"], + timestamp=timestamp, + action_type=raw["action_type"], + target=raw["target"], + change=change, + source=raw.get("source", self.source), + raw_ref=raw.get("raw_ref"), + ) + except KeyError as exc: + raise AdapterError(f"Generic record missing required field: {exc}") from exc + except (TypeError, ValueError) as exc: + raise AdapterError(f"Invalid generic record: {exc}") from exc diff --git a/examples/agent-action-monitor/src/dusk/actions/analyse.py b/examples/agent-action-monitor/src/dusk/actions/analyse.py new file mode 100644 index 000000000..1cf675575 --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/analyse.py @@ -0,0 +1,387 @@ +"""Score agent behavior against its baseline and explain each signal.""" + +from __future__ import annotations + +import logging +import math +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any + +from dusk.actions.baseline import Baseline, action_features +from dusk.actions.event import AgentAction +from dusk.trace.vector import sie_extract, sie_score + +if TYPE_CHECKING: + from dusk.actions.offense_memory import OffenseRecord + from dusk.config import Config + +logger = logging.getLogger("dusk.actions.analyse") + +#: Weight each novelty signal contributes to the anomaly score. +_W_NEW_ACTION_TYPE = 0.55 +_W_NEW_TARGET_CLASS = 0.4 +_W_NEW_TOKENS = 0.2 +_W_NEW_CHANGE_VALUES = 0.25 +_W_UNKNOWN_AGENT = 0.5 +_W_SENSITIVE = 0.35 +#: Extra contribution when SIE's reranker finds no close match in the +#: agent's raw history, even though the deterministic feature checks above +#: found nothing new. Only fires when SIE is configured and reachable; the +#: rule-based score above is unchanged otherwise. +_W_LOW_SEMANTIC_SIMILARITY = 0.2 +#: Rerank score below this is "no close match". sie_score() bounds the raw logit into [0, 1] via +#: sigmoid, but this floor is a heuristic, not an empirically calibrated cutoff. +_SEMANTIC_SIMILARITY_FLOOR = 0.3 +#: Extra contribution when SIE's zero-shot extractor (GLiNER) surfaces a +#: privileged term the static frozenset below doesn't already cover. Slightly +#: below _W_SENSITIVE since it's a probabilistic match, not an exact one. +_W_EXTRACTED_SENSITIVE = 0.3 +#: Below this confidence, a zero-shot extraction is treated as noise, not evidence. +_EXTRACT_CONFIDENCE_FLOOR = 0.5 +#: Only these GLiNER labels indicate privilege escalation; "resource"/"segment"/"port" (also in +#: DEFAULT_EXTRACT_LABELS) are neutral descriptors and shouldn't add score on their own. +_PRIVILEGED_LABELS = frozenset({"role", "privilege"}) + +#: MITRE ATT&CK technique per normalised action type. +_ATTCK: dict[str, str] = { + "firewall_rule_change": "T1562.004 Impair Defenses: Disable or Modify System Firewall", + "port_change": "T1562.004 Impair Defenses: Disable or Modify System Firewall", + "route_change": "T1599 Network Boundary Bridging", + "segment_change": "T1599 Network Boundary Bridging", + "role_assignment": "T1098 Account Manipulation", + "unknown": "T1078 Valid Accounts", +} + +#: MITRE ATLAS technique describing the agent-level cause. +_ATLAS = "AML.T0051 LLM Prompt Injection" + +#: Sensitive change values that signal privilege escalation when newly introduced. +_SENSITIVE_VALUES = frozenset({"owner", "admin", "root", "0.0.0.0/0"}) +#: Target tokens that indicate a cross-boundary or restricted reach. +_SENSITIVE_TOKENS = frozenset({"restricted", "guest", "self", "owner", "global", "all"}) +#: The union, used to test whether a term is sensitive. +_SENSITIVE = _SENSITIVE_VALUES | _SENSITIVE_TOKENS + + +@dataclass +class AnalysisResult: + """The outcome of analysing one action against a baseline. + + Attributes: + agent_id: The acting agent. + action_type: The action's normalised verb. + target: What was acted on. + score: Anomaly score in ``0..1``; higher is more anomalous. + reasons: Human-readable explanations of the score. + mitre_attack: The mapped MITRE ATT&CK technique. + mitre_atlas: The mapped MITRE ATLAS technique. + blast_radius: Coarse impact estimate, ``"low"``, ``"medium"`` or + ``"high"``. + predicted_next: What an attacker would likely do next. + """ + + agent_id: str + action_type: str + target: str + score: float + reasons: list[str] = field(default_factory=list) + mitre_attack: str = "" + mitre_atlas: str = "" + blast_radius: str = "low" + predicted_next: str = "" + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serialisable representation of the result.""" + return { + "agent_id": self.agent_id, + "action_type": self.action_type, + "target": self.target, + "score": round(self.score, 4), + "reasons": self.reasons, + "mitre_attack": self.mitre_attack, + "mitre_atlas": self.mitre_atlas, + "blast_radius": self.blast_radius, + "predicted_next": self.predicted_next, + } + + +def _blast_radius(action: AgentAction, features: dict[str, Any]) -> str: + """Estimate how much damage the action could do.""" + sensitive = (features["tokens"] | features["change_values"]) & ( + _SENSITIVE_TOKENS | _SENSITIVE_VALUES + ) + if sensitive: + return "high" + if action.action_type in ("firewall_rule_change", "role_assignment", "route_change"): + return "medium" + return "low" + + +def _predicted_next(action: AgentAction) -> str: + """Predict the attacker's likely next move after this action.""" + mapping = { + "firewall_rule_change": ( + "expect lateral movement into the newly reachable segment; watch for " + "east-west connections from this agent" + ), + "route_change": ( + "expect traffic redirection or interception; watch for new flows toward " + "the changed next hop" + ), + "segment_change": ( + "expect access into the resized or new segment; watch for first-time connections there" + ), + "role_assignment": ( + "expect privilege use; watch for actions that the newly granted role " + "permits but this agent never took before" + ), + } + return mapping.get( + action.action_type, + "watch this agent for further actions outside its established pattern", + ) + + +def _extracted_sensitive_terms(features: dict[str, Any]) -> set[str]: + """Pull privileged terms from the action's target/change text via SIE extract. + + Returns an empty set whenever SIE is not configured/reachable, or when + every extraction is low-confidence or labeled as a neutral descriptor + rather than role/privilege (see _EXTRACT_CONFIDENCE_FLOOR, _PRIVILEGED_LABELS). + """ + text = " ".join(sorted(features["tokens"] | features["change_values"])) + if not text: + return set() + return { + term.text.lower() + for term in sie_extract(text) + if term.score >= _EXTRACT_CONFIDENCE_FLOOR and term.label.lower() in _PRIVILEGED_LABELS + } + + +def _semantic_novelty( + action: AgentAction, agent_history: list[AgentAction] +) -> tuple[float, str | None]: + """Rerank ``action`` against the agent's raw history via SIE's cross-encoder. + + Returns ``(0.0, None)`` whenever there is no history to compare against + or SIE is not configured/reachable, so the deterministic score above is + unchanged in the default (no-SIE) case. + """ + if not agent_history: + return 0.0, None + + query = f"{action.action_type} {action.target}" + candidates = [f"{a.action_type} {a.target}" for a in agent_history] + scores = sie_score(query, candidates) + if not scores: + return 0.0, None + + best = max(scores) + if best < _SEMANTIC_SIMILARITY_FLOOR: + return ( + _W_LOW_SEMANTIC_SIMILARITY, + f"SIE rerank finds no close match in this agent's recorded history (best={best:.2f})", + ) + return 0.0, None + + +def _extra_sie_signals( + action: AgentAction, features: dict[str, Any], agent_history: list[AgentAction] +) -> tuple[float, list[str]]: + """Optional SIE-backed signals: extracted privileged terms and rerank novelty. + + Both are no-ops (contribute nothing) when SIE is not configured/reachable, + so the deterministic score above is unchanged in the default case. + """ + extra_score = 0.0 + reasons: list[str] = [] + + extracted = _extracted_sensitive_terms(features) - _SENSITIVE + if extracted: + extra_score += _W_EXTRACTED_SENSITIVE + reasons.append(f"SIE extract flags additional privileged terms {sorted(extracted)}") + + novelty_score, novelty_reason = _semantic_novelty(action, agent_history) + if novelty_reason: + extra_score += novelty_score + reasons.append(novelty_reason) + + return extra_score, reasons + + +def _offense_similarity( + action: AgentAction, features: dict[str, Any], offense: OffenseRecord +) -> float: + """How closely ``action`` matches a single past offense, in ``0..1``. + + Requires the same action type to count at all -- a firewall change + doesn't make a later role assignment suspicious just because both were + once blocked. Beyond that, target class and shared tokens each add + partial credit, so an exact repeat of the same target scores highest + while a same-type action against a related-but-different target still + contributes something. + """ + if action.action_type != offense.action_type: + return 0.0 + similarity = 0.5 + if features["target_class"] and features["target_class"] == offense.target_class: + similarity += 0.3 + if features["tokens"] & set(offense.tokens): + similarity += 0.2 + return min(1.0, similarity) + + +def _decay(offense: OffenseRecord, half_life_days: float) -> float: + """Exponential decay: 1.0 for a brand-new offense, 0.5 at one half-life, etc.""" + age_days = max(0.0, (datetime.now(UTC) - offense.timestamp).total_seconds() / 86400) + return float(math.pow(0.5, age_days / half_life_days)) + + +def _repeat_offense_signal( + action: AgentAction, + features: dict[str, Any], + offenses: list[OffenseRecord] | None, + config: Config | None, +) -> tuple[float, str | None]: + """Score how much ``action`` resembles this agent's own past refusals. + + Takes the single best-matching prior offense rather than summing across + all of them, so an attacker cannot inflate the contribution by + triggering many weak matches -- only the strongest, most relevant prior + refusal counts, and its weight decays with age. Returns ``(0.0, None)`` + when there is no meaningful match, so a clean-history agent is + completely unaffected by this signal. + """ + if not offenses: + return 0.0, None + + from dusk.config import get_config + + cfg = config or get_config() + best_offense: OffenseRecord | None = None + best_weight = 0.0 + for offense in offenses: + similarity = _offense_similarity(action, features, offense) + if similarity <= 0.0: + continue + weight = similarity * _decay(offense, cfg.repeat_offense_half_life_days) + if weight > best_weight: + best_weight = weight + best_offense = offense + + if best_offense is None or best_weight <= 0.0: + return 0.0, None + + contribution = min(cfg.repeat_offense_max_contribution, best_weight) + reason = ( + f"resembles a prior {best_offense.verdict} action by this agent " + f"(trace {best_offense.trace_id}, {best_offense.timestamp.date().isoformat()})" + ) + return contribution, reason + + +def analyse( + baseline: Baseline, + action: AgentAction, + agent_history: list[AgentAction] | None = None, + offenses: list[OffenseRecord] | None = None, + config: Config | None = None, +) -> AnalysisResult: + """Score and explain ``action`` against ``baseline``. + + Args: + baseline: The learned per-agent baseline. + action: The action to evaluate. + agent_history: The agent's raw known-good actions, used for an + optional SIE-reranked semantic novelty check on top of the + deterministic feature checks below. Omit or pass an empty list + to skip this signal entirely. + offenses: The agent's past refused verdicts, used for the + repeat-offense signal. Omit or pass an empty list to skip this + signal entirely -- a clean-history agent is unaffected. + config: Configuration providing ``repeat_offense_max_contribution`` + and ``repeat_offense_half_life_days``. Defaults to the + process-wide singleton; only consulted when ``offenses`` is + non-empty. + + Returns: + An :class:`AnalysisResult` with score, reasons, and mappings. + """ + features = action_features(action) + profile = baseline.profile_for(action.agent_id) + reasons: list[str] = [] + score = 0.0 + + if profile is None or profile.count == 0: + score += _W_UNKNOWN_AGENT + reasons.append( + f"agent '{action.agent_id}' has no established baseline; " + f"its behaviour cannot be vouched for" + ) + sensitive = (features["tokens"] | features["change_values"]) & _SENSITIVE + if sensitive: + score += _W_SENSITIVE + reasons.append(f"action touches sensitive terms {sorted(sensitive)}") + else: + if features["action_type"] not in profile.action_types: + score += _W_NEW_ACTION_TYPE + reasons.append( + f"action type '{features['action_type']}' is new for this agent, " + f"which normally does {sorted(profile.action_types)}" + ) + if features["target_class"] and features["target_class"] not in profile.target_classes: + score += _W_NEW_TARGET_CLASS + reasons.append( + f"target class '{features['target_class']}' is new for this agent, " + f"which normally touches {sorted(profile.target_classes)}" + ) + new_tokens = features["tokens"] - profile.tokens + if new_tokens: + score += _W_NEW_TOKENS + reasons.append(f"target introduces unseen terms {sorted(new_tokens)}") + new_values = features["change_values"] - profile.change_values + if new_values: + score += _W_NEW_CHANGE_VALUES + reasons.append(f"change introduces unseen values {sorted(new_values)}") + sensitive_new = (new_tokens | new_values) & _SENSITIVE + if sensitive_new: + score += _W_SENSITIVE + reasons.append( + f"newly introduces sensitive or privileged terms {sorted(sensitive_new)}" + ) + + extra_score, extra_reasons = _extra_sie_signals(action, features, agent_history or []) + score += extra_score + reasons.extend(extra_reasons) + + repeat_score, repeat_reason = _repeat_offense_signal(action, features, offenses, config) + if repeat_reason: + score += repeat_score + reasons.append(repeat_reason) + + score = min(1.0, score) + blast = _blast_radius(action, features) + if not reasons: + reasons.append("action matches the agent's established pattern") + + result = AnalysisResult( + agent_id=action.agent_id, + action_type=action.action_type, + target=action.target, + score=score, + reasons=reasons, + mitre_attack=_ATTCK.get(action.action_type, _ATTCK["unknown"]), + mitre_atlas=_ATLAS, + blast_radius=blast, + predicted_next=_predicted_next(action), + ) + logger.debug( + "analysed agent=%s action_type=%s score=%.2f blast=%s", + action.agent_id, + action.action_type, + score, + blast, + ) + return result diff --git a/examples/agent-action-monitor/src/dusk/actions/baseline.py b/examples/agent-action-monitor/src/dusk/actions/baseline.py new file mode 100644 index 000000000..66760f4f0 --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/baseline.py @@ -0,0 +1,119 @@ +"""Deterministic per-agent behavioral baselines and feature extraction.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from typing import Any + +from dusk.actions.event import AgentAction + +#: Split a target identifier into lowercase word tokens. +_TOKEN_RE = re.compile(r"[a-z0-9]+") + + +def target_class(target: str) -> str: + """Return the coarse class of a target, its first token. + + For example ``"fw-corp-https"`` and ``"fw-guest-to-restricted"`` share the + class ``"fw"``, while ``"seg-corporate"`` is ``"seg"``. The class captures + the kind of resource an agent touches without overfitting to exact names. + """ + tokens = _TOKEN_RE.findall(target.lower()) + return tokens[0] if tokens else "" + + +def target_tokens(target: str) -> set[str]: + """Return the set of lowercase word tokens in a target identifier.""" + return set(_TOKEN_RE.findall(target.lower())) + + +def _flatten_scalars(payload: Any, values: set[str], *, _depth: int = 0) -> None: # noqa: ANN401 + """Collect nested scalar values with a defensive recursion limit.""" + if _depth > 10: + return + if isinstance(payload, dict): + for value in payload.values(): + _flatten_scalars(value, values, _depth=_depth + 1) + elif isinstance(payload, list): + for value in payload: + _flatten_scalars(value, values, _depth=_depth + 1) + elif isinstance(payload, (str, int, float, bool)): + values.add(str(payload).lower()) + + +def _change_values(change: dict[str, Any]) -> set[str]: + """Flatten a change delta into a set of stringified scalar values, at any nesting depth.""" + values: set[str] = set() + for side in ("before", "after"): + _flatten_scalars(change.get(side), values) + return values + + +@dataclass +class AgentProfile: + """The learned normal behaviour of a single agent.""" + + agent_id: str + action_types: set[str] = field(default_factory=set) + target_classes: set[str] = field(default_factory=set) + tokens: set[str] = field(default_factory=set) + change_values: set[str] = field(default_factory=set) + count: int = 0 + + def observe(self, action: AgentAction) -> None: + """Fold one known-good action into the profile.""" + self.action_types.add(action.action_type) + self.target_classes.add(target_class(action.target)) + self.tokens |= target_tokens(action.target) + self.change_values |= _change_values(action.change) + self.count += 1 + + +def action_features(action: AgentAction) -> dict[str, Any]: + """Extract the comparable features of an action. + + Returns the action's action type, target class, target tokens, and change + values. This is the single seam a vector backend would replace. + """ + return { + "action_type": action.action_type, + "target_class": target_class(action.target), + "tokens": target_tokens(action.target), + "change_values": _change_values(action.change), + } + + +class Baseline: + """A collection of per-agent profiles learned from known-good actions.""" + + def __init__(self) -> None: + self._profiles: dict[str, AgentProfile] = {} + + @classmethod + def learn(cls, actions: list[AgentAction]) -> Baseline: + """Build a baseline from a history of known-good actions.""" + baseline = cls() + for action in actions: + baseline.observe(action) + return baseline + + def observe(self, action: AgentAction) -> None: + """Fold one known-good action into its agent's profile.""" + profile = self._profiles.get(action.agent_id) + if profile is None: + profile = AgentProfile(agent_id=action.agent_id) + self._profiles[action.agent_id] = profile + profile.observe(action) + + def profile_for(self, agent_id: str) -> AgentProfile | None: + """Return the profile for an agent, or ``None`` if unseen.""" + return self._profiles.get(agent_id) + + @property + def agents(self) -> list[str]: + """The sorted list of agents the baseline has profiles for.""" + return sorted(self._profiles) + + def __len__(self) -> int: + return len(self._profiles) diff --git a/examples/agent-action-monitor/src/dusk/actions/event.py b/examples/agent-action-monitor/src/dusk/actions/event.py new file mode 100644 index 000000000..18329123b --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/event.py @@ -0,0 +1,125 @@ +"""The AgentAction event: one normalised control-plane action. + +An :class:`AgentAction` is controller-agnostic by design. Whatever system an +agent used to change the network (a cloud API, an SDN controller, a policy +endpoint), the event records the same canonical facts: who acted, when, what +kind of change, on what target, and the before/after delta. Source-specific +shapes are mapped onto this schema by adapters. + +This object describes what happened, not how bad it is. It carries no +severity, score, blast radius, or verdict; those belong to later layers. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime +from typing import Any + +#: The normalised set of action verbs. Anything an adapter cannot classify +#: maps to ``"unknown"`` rather than being invented. +KNOWN_ACTION_TYPES: frozenset[str] = frozenset( + { + "firewall_rule_change", + "route_change", + "segment_change", + "role_assignment", + "port_change", + "unknown", + } +) + + +@dataclass(frozen=True) +class AgentAction: + """One normalised action an agent performed on the network control plane. + + Attributes: + agent_id: Identity of the acting agent (service account, role, or + agent name as the controller reports it). Must be non-empty. + timestamp: When the action occurred. Must be timezone-aware. + action_type: Normalised verb, one of :data:`KNOWN_ACTION_TYPES`. + target: What was acted on (resource id, rule name, segment). Must be + non-empty. + change: Structured delta with keys ``"before"`` and ``"after"``; + either may be ``None`` for a create or delete. + source: Originating controller, for example ``"azure"`` or + ``"generic"``. + raw_ref: Opaque reference back to the original record (an id), never + the full payload. May be ``None``. + """ + + agent_id: str + timestamp: datetime + action_type: str + target: str + change: dict[str, Any] + source: str + raw_ref: str | None = None + + def __post_init__(self) -> None: + """Validate the event, raising ValueError on invalid input. + + Raises: + ValueError: If ``agent_id`` or ``target`` is empty, ``timestamp`` + is not timezone-aware, ``action_type`` is not a known verb, or + ``change`` is not a mapping. Values are never silently coerced. + """ + if not isinstance(self.agent_id, str) or not self.agent_id.strip(): + raise ValueError("agent_id must be a non-empty string") + if not isinstance(self.target, str) or not self.target.strip(): + raise ValueError("target must be a non-empty string") + if not isinstance(self.timestamp, datetime): + raise ValueError("timestamp must be a datetime") + if self.timestamp.tzinfo is None or self.timestamp.utcoffset() is None: + raise ValueError("timestamp must be timezone-aware") + if self.action_type not in KNOWN_ACTION_TYPES: + raise ValueError( + f"action_type '{self.action_type}' is not one of the known " + f"types: {', '.join(sorted(KNOWN_ACTION_TYPES))}" + ) + if not isinstance(self.change, dict): + raise ValueError("change must be a dict with 'before' and 'after' keys") + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-safe dict, with the timestamp as an ISO 8601 string.""" + return { + "agent_id": self.agent_id, + "timestamp": self.timestamp.isoformat(), + "action_type": self.action_type, + "target": self.target, + "change": self.change, + "source": self.source, + "raw_ref": self.raw_ref, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> AgentAction: + """Build an :class:`AgentAction` from a :meth:`to_dict` mapping. + + Args: + data: A mapping as produced by :meth:`to_dict`, with the timestamp + as an ISO 8601 string. + + Returns: + The reconstructed :class:`AgentAction`. + + Raises: + ValueError: If a required key is missing or a value is invalid. + """ + try: + timestamp = datetime.fromisoformat(data["timestamp"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError(f"Invalid or missing timestamp: {exc}") from exc + try: + return cls( + agent_id=data["agent_id"], + timestamp=timestamp, + action_type=data["action_type"], + target=data["target"], + change=data["change"], + source=data["source"], + raw_ref=data.get("raw_ref"), + ) + except KeyError as exc: + raise ValueError(f"Missing required field: {exc}") from exc diff --git a/examples/agent-action-monitor/src/dusk/actions/ingest.py b/examples/agent-action-monitor/src/dusk/actions/ingest.py new file mode 100644 index 000000000..a28a941ea --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/ingest.py @@ -0,0 +1,72 @@ +"""Normalize JSON source records into AgentAction events.""" + +from __future__ import annotations + +import json +import logging +import os + +from dusk.actions.event import AgentAction + +logger = logging.getLogger("dusk.actions.ingest") + + +def ingest_file(path: str, source: str) -> list[AgentAction]: + """Read a JSON list of raw ``source`` records into AgentAction events. + + Args: + path: Filesystem path to a JSON file containing a list of records. + source: Canonical source name selecting the adapter (for example + ``"azure"`` or ``"generic"``). + + Returns: + The list of successfully normalised :class:`AgentAction` events. + + Raises: + FileNotFoundError: If ``path`` does not exist. + ValueError: If the file is empty, not valid JSON, or its top-level + JSON is not a list. + """ + # Imported here so the adapter registry loads lazily and ingest stays free + # of an import cycle (ingest -> normaliser -> adapters). + from dusk.actions.adapters.base import AdapterError + from dusk.actions.normaliser import normalise_record + + if not os.path.exists(path): + logger.critical("Action file not found: %s", path) + raise FileNotFoundError(f"Action file not found: {path}") + + if os.path.getsize(path) == 0: + logger.warning("Action file is empty: %s", path) + raise ValueError(f"Action file is empty: {path}") + + with open(path, encoding="utf-8") as handle: + try: + payload = json.load(handle) + except json.JSONDecodeError as exc: + logger.critical("Action file is not valid JSON: %s: %s", path, exc) + raise ValueError(f"Action file is not valid JSON: {path}: {exc}") from exc + + if not isinstance(payload, list): + raise ValueError( + f"Action file must contain a JSON list of records, got {type(payload).__name__}: {path}" + ) + + actions: list[AgentAction] = [] + for index, record in enumerate(payload): + if not isinstance(record, dict): + logger.warning("Skipping record %d in %s: not a JSON object", index, path) + continue + try: + actions.append(normalise_record(source, record)) + except AdapterError as exc: + logger.warning("Skipping record %d in %s: %s", index, path, exc) + + logger.info( + "Ingested %d action(s) from %s via '%s' adapter (%d record(s) skipped)", + len(actions), + path, + source, + len(payload) - len(actions), + ) + return actions diff --git a/examples/agent-action-monitor/src/dusk/actions/normaliser.py b/examples/agent-action-monitor/src/dusk/actions/normaliser.py new file mode 100644 index 000000000..09d6cf730 --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/normaliser.py @@ -0,0 +1,64 @@ +"""Route source records through registered AgentAction adapters.""" + +from __future__ import annotations + +import logging +from typing import Any + +from dusk.actions.adapters.azure import AzureAdapter +from dusk.actions.adapters.base import AdapterError, SourceAdapter +from dusk.actions.adapters.bedrock import BedrockAdapter +from dusk.actions.adapters.generic import GenericAdapter +from dusk.actions.event import AgentAction + +logger = logging.getLogger("dusk.actions.normaliser") + +#: Registry of source adapters, keyed by canonical source name. +_REGISTRY: dict[str, SourceAdapter] = { + GenericAdapter.source: GenericAdapter(), + AzureAdapter.source: AzureAdapter(), + BedrockAdapter.source: BedrockAdapter(), +} + + +def known_sources() -> list[str]: + """Return the sorted list of registered source names.""" + return sorted(_REGISTRY) + + +def get_adapter(source: str) -> SourceAdapter: + """Return the adapter registered for ``source``. + + Args: + source: Canonical source name, for example ``"azure"`` or + ``"generic"``. + + Returns: + The registered :class:`SourceAdapter`. + + Raises: + AdapterError: If no adapter is registered for ``source``. + """ + adapter = _REGISTRY.get(source) + if adapter is None: + raise AdapterError( + f"No adapter registered for source '{source}'. Known sources: " + f"{', '.join(known_sources())}" + ) + return adapter + + +def normalise_record(source: str, raw: dict[str, Any]) -> AgentAction: + """Normalise one raw record from ``source`` into an :class:`AgentAction`. + + Args: + source: Canonical source name selecting the adapter. + raw: One raw record in that source's native shape. + + Returns: + The canonical :class:`AgentAction`. + + Raises: + AdapterError: If the source is unknown or the record is malformed. + """ + return get_adapter(source).parse(raw) diff --git a/examples/agent-action-monitor/src/dusk/actions/offense_memory.py b/examples/agent-action-monitor/src/dusk/actions/offense_memory.py new file mode 100644 index 000000000..0f162482f --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/offense_memory.py @@ -0,0 +1,261 @@ +"""Bounded, durable history of refused actions, partitioned by agent.""" + +from __future__ import annotations + +import json +import logging +import threading +from collections import OrderedDict +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import asdict, dataclass, field +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +logger = logging.getLogger("dusk.actions.offense_memory") + +#: Retention limits bound both memory use and the persisted file. +_MAX_OFFENSES_PER_AGENT = 50 +_MAX_TRACKED_AGENTS = 500 + + +@dataclass(frozen=True) +class OffenseRecord: + """One persisted refusal: enough to cite it later and judge similarity to it. + + Attributes: + trace_id: The gate's trace_id for the refused verdict, so a later + citation ("repeat of trace ") is checkable against real logs. + agent_id: The agent this offense belongs to. + action_type: The refused action's normalised verb. + target_class: The refused action's coarse target class (see + :func:`dusk.actions.baseline.target_class`). + tokens: The refused action's target tokens, for similarity matching + against a new action. + verdict: The verdict rendered, ``WOULD-BLOCK`` or ``BLOCK``. + timestamp: When the offense was recorded, UTC, timezone-aware. + """ + + trace_id: str + agent_id: str + action_type: str + target_class: str + tokens: tuple[str, ...] + verdict: str + timestamp: datetime + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-safe representation.""" + data = asdict(self) + data["timestamp"] = self.timestamp.isoformat() + data["tokens"] = list(self.tokens) + return data + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> OffenseRecord: + """Reconstruct a record from :meth:`to_dict` output. + + Raises: + ValueError: If a required key is missing or malformed. + """ + try: + return cls( + trace_id=data["trace_id"], + agent_id=data["agent_id"], + action_type=data["action_type"], + target_class=data["target_class"], + tokens=tuple(data["tokens"]), + verdict=data["verdict"], + timestamp=datetime.fromisoformat(data["timestamp"]), + ) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError(f"Invalid offense record: {exc}") from exc + + +@dataclass +class _AgentOffenses: + records: list[OffenseRecord] = field(default_factory=list) + + +class OffenseMemory: + """Thread-safe offense store with coalesced, single-writer persistence. + + The dirty flag prevents request bursts from creating an unbounded write + queue. All scheduling state shares the data lock, preserving the + single-writer invariant during concurrent first use. + """ + + def __init__(self, storage_path: str | None = None) -> None: + """Create the store, loading any existing state from ``storage_path``. + + Args: + storage_path: File to persist offenses to. When ``None``, the + store is in-memory only for the life of this instance (used + by tests that don't care about durability). + """ + self._storage_path = Path(storage_path) if storage_path else None + self._lock = threading.Lock() + self._by_agent: OrderedDict[str, _AgentOffenses] = OrderedDict() + if self._storage_path is not None: + self._load(self._storage_path) + self._executor: ThreadPoolExecutor | None = None + self._write_in_flight = False + self._dirty = False + self._last_write: Future[None] | None = None + self._last_persist_error: str | None = None + self._closed = False + + def _touch(self, agent_id: str) -> _AgentOffenses: + """Return the agent record and enforce LRU retention. Lock required.""" + offenses = self._by_agent.get(agent_id) + if offenses is None: + offenses = _AgentOffenses() + self._by_agent[agent_id] = offenses + else: + self._by_agent.move_to_end(agent_id) + while len(self._by_agent) > _MAX_TRACKED_AGENTS: + evicted_id, _ = self._by_agent.popitem(last=False) + logger.info("Evicted offense memory for agent %s (tracked-agent cap)", evicted_id) + return offenses + + def _load(self, storage_path: Path) -> None: + if not storage_path.exists(): + return + try: + with storage_path.open(encoding="utf-8") as fh: + raw = json.load(fh) + except (OSError, json.JSONDecodeError) as exc: + logger.warning( + "Could not read offense memory from %s, starting empty: %s", + storage_path, + exc, + ) + return + if not isinstance(raw, dict): + logger.warning( + "Offense memory file %s did not contain a mapping, starting empty", + storage_path, + ) + return + for agent_id, entries in raw.items(): + if not isinstance(entries, list): + continue + records: list[OffenseRecord] = [] + for entry in entries: + try: + records.append(OffenseRecord.from_dict(entry)) + except ValueError as exc: + logger.warning("Skipping malformed offense record: %s", exc) + if records: + capped = records[-_MAX_OFFENSES_PER_AGENT:] + self._by_agent[agent_id] = _AgentOffenses(records=capped) + while len(self._by_agent) > _MAX_TRACKED_AGENTS: + self._by_agent.popitem(last=False) + logger.info( + "Loaded offense memory for %d agent(s) from %s", + len(self._by_agent), + storage_path, + ) + + def _write_to_disk(self, payload: dict[str, Any]) -> None: + if self._storage_path is None: + return + try: + self._storage_path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = self._storage_path.with_suffix(".tmp") + with tmp_path.open("w", encoding="utf-8") as fh: + json.dump(payload, fh, indent=2) + tmp_path.replace(self._storage_path) + self._last_persist_error = None + except OSError as exc: + self._last_persist_error = str(exc) + logger.warning("Could not persist offense memory to %s: %s", self._storage_path, exc) + + def _persist_loop(self) -> None: + """Persist snapshots until no changes occurred during the last write.""" + while True: + with self._lock: + self._dirty = False + payload = { + agent_id: [r.to_dict() for r in offenses.records] + for agent_id, offenses in self._by_agent.items() + } + self._write_to_disk(payload) + with self._lock: + if not self._dirty: + self._write_in_flight = False + return + + def _schedule_persist(self) -> None: + """Mark state dirty and start the writer if needed. Lock required.""" + if self._storage_path is None or self._closed: + return + self._dirty = True + if self._write_in_flight: + return + self._write_in_flight = True + if self._executor is None: + self._executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="offense-memory-writer" + ) + self._last_write = self._executor.submit(self._persist_loop) + + def flush(self, timeout: float | None = None) -> None: + """Block until every write scheduled before this call is durable.""" + with self._lock: + future = self._last_write + if future is not None: + future.result(timeout=timeout) + + def close(self, timeout: float | None = None) -> None: + """Flush and stop the writer. Safe to call more than once.""" + self.flush(timeout=timeout) + with self._lock: + self._closed = True + executor = self._executor + self._executor = None + if executor is not None: + executor.shutdown(wait=True) + + @property + def last_persist_error(self) -> str | None: + """Return the latest disk-write error for health reporting.""" + return self._last_persist_error + + def record( + self, + trace_id: str, + agent_id: str, + action_type: str, + target_class: str, + tokens: set[str], + verdict: str, + ) -> None: + """Record a refused verdict for ``agent_id`` and schedule a write.""" + entry = OffenseRecord( + trace_id=trace_id, + agent_id=agent_id, + action_type=action_type, + target_class=target_class, + tokens=tuple(sorted(tokens)), + verdict=verdict, + timestamp=datetime.now(UTC), + ) + with self._lock: + offenses = self._touch(agent_id) + offenses.records.append(entry) + if len(offenses.records) > _MAX_OFFENSES_PER_AGENT: + del offenses.records[: len(offenses.records) - _MAX_OFFENSES_PER_AGENT] + self._schedule_persist() + + def offenses_for(self, agent_id: str) -> list[OffenseRecord]: + """Return the recorded offenses for one agent, oldest first.""" + with self._lock: + offenses = self._by_agent.get(agent_id) + return list(offenses.records) if offenses else [] + + def clear(self) -> None: + """Wipe all recorded offenses. Test-only / operator reset hook.""" + with self._lock: + self._by_agent.clear() + self._schedule_persist() diff --git a/examples/agent-action-monitor/src/dusk/actions/verdict.py b/examples/agent-action-monitor/src/dusk/actions/verdict.py new file mode 100644 index 000000000..f3b790130 --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/actions/verdict.py @@ -0,0 +1,136 @@ +"""Render ALLOW, WOULD-BLOCK, or BLOCK from behavioral analysis.""" + +from __future__ import annotations + +import logging +import uuid +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +from dusk.actions.analyse import AnalysisResult, analyse +from dusk.actions.baseline import Baseline, action_features +from dusk.actions.event import AgentAction +from dusk.config import Config, get_config + +if TYPE_CHECKING: + from dusk.actions.offense_memory import OffenseMemory + +logger = logging.getLogger("dusk.actions.verdict") + +#: Verdict strings. +ALLOW = "ALLOW" +WOULD_BLOCK = "WOULD-BLOCK" +BLOCK = "BLOCK" + + +@dataclass +class GateVerdict: + """A gate decision about a single action. + + Attributes: + verdict: One of ``ALLOW``, ``WOULD-BLOCK``, ``BLOCK``. + analysis: The :class:`AnalysisResult` the verdict is based on. + trace_id: Unique id for this decision, generated once here so every + downstream consumer (the HTTP response, offense memory, n8n + webhooks) cites the same identifier for the same decision. + """ + + verdict: str + analysis: AnalysisResult + trace_id: str = field(default_factory=lambda: uuid.uuid4().hex) + + @property + def refused(self) -> bool: + """True when the action was not allowed.""" + return self.verdict != ALLOW + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serialisable representation of the verdict.""" + return {"trace_id": self.trace_id, "verdict": self.verdict, **self.analysis.to_dict()} + + +class ActionGate: + """Inline gate: learn normal behaviour, then judge new actions.""" + + def __init__( + self, + baseline: Baseline | None = None, + config: Config | None = None, + enforce: bool = False, + offense_memory: OffenseMemory | None = None, + ) -> None: + """Create the gate. + + Args: + baseline: A pre-learned baseline, or ``None`` to start empty and + learn with :meth:`learn`. + config: Configuration providing ``gate_block_threshold``. Defaults + to the process-wide singleton. + enforce: When ``True``, refused actions are BLOCK; otherwise they + are WOULD-BLOCK (watch mode). + offense_memory: Persisted per-agent record of past refusals, used + for the repeat-offense scoring signal. When ``None``, that + signal is skipped entirely -- every evaluation behaves as it + did before this store existed. + """ + self.config = config if config is not None else get_config() + self.baseline = baseline if baseline is not None else Baseline() + self.enforce = enforce + self.offense_memory = offense_memory + self._history: dict[str, list[AgentAction]] = {} + + def learn(self, actions: list[AgentAction]) -> None: + """Fold known-good actions into the baseline.""" + for action in actions: + self.baseline.observe(action) + self._history.setdefault(action.agent_id, []).append(action) + logger.info( + "gate baseline learned: %d agent(s) from %d action(s)", + len(self.baseline), + len(actions), + ) + + def evaluate(self, action: AgentAction) -> GateVerdict: + """Analyse one action and render a verdict.""" + offenses = ( + self.offense_memory.offenses_for(action.agent_id) if self.offense_memory else None + ) + result = analyse( + self.baseline, + action, + agent_history=self._history.get(action.agent_id), + offenses=offenses, + config=self.config, + ) + if result.score >= self.config.gate_block_threshold: + verdict = BLOCK if self.enforce else WOULD_BLOCK + logger.error( + "gate refused: verdict=%s agent=%s action_type=%s target=%s " + "score=%.2f attack=%s atlas=%s blast=%s", + verdict, + result.agent_id, + result.action_type, + result.target, + result.score, + result.mitre_attack, + result.mitre_atlas, + result.blast_radius, + ) + else: + verdict = ALLOW + gate_verdict = GateVerdict(verdict=verdict, analysis=result) + if gate_verdict.refused and self.offense_memory is not None: + features = action_features(action) + self.offense_memory.record( + trace_id=gate_verdict.trace_id, + agent_id=action.agent_id, + action_type=action.action_type, + target_class=features["target_class"], + tokens=features["tokens"], + verdict=verdict, + ) + return gate_verdict + + def evaluate_all(self, actions: list[AgentAction]) -> list[GateVerdict]: + """Evaluate a sequence of actions in order.""" + return [self.evaluate(action) for action in actions] diff --git a/examples/agent-action-monitor/src/dusk/api.py b/examples/agent-action-monitor/src/dusk/api.py new file mode 100644 index 000000000..b8ce43bad --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/api.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +import logging +import os +import threading +from typing import TYPE_CHECKING + +from dotenv import load_dotenv +from flask import Flask, jsonify, request +from flask_cors import CORS + +if TYPE_CHECKING: + from dusk.actions.verdict import ActionGate + from dusk.trace.models import TraceDecision + +load_dotenv() + +app = Flask(__name__) +CORS(app) +# Bound public input without constraining normal AgentAction payloads. +app.config["MAX_CONTENT_LENGTH"] = 1 * 1024 * 1024 + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +_gate_engine: ActionGate | None = None +_gate_lock = threading.Lock() + +#: Baseline failure exposed by the health endpoint. +_baseline_load_error: str | None = None + + +def _load_gate_engine() -> ActionGate: + from dusk.actions.ingest import ingest_file + from dusk.actions.offense_memory import OffenseMemory + from dusk.actions.verdict import ActionGate + from dusk.config import get_config + + global _baseline_load_error + _baseline_load_error = None + + baseline_path = os.getenv("DUSK_GATE_BASELINE_PATH", "") + baseline_source = os.getenv("DUSK_GATE_BASELINE_SOURCE", "generic") + + config = get_config() + offense_memory = OffenseMemory(storage_path=config.offense_memory_path or None) + gate_engine = ActionGate(config=config, enforce=config.enforce, offense_memory=offense_memory) + if baseline_path: + try: + known_good = ingest_file(baseline_path, baseline_source) + gate_engine.learn(known_good) + except (FileNotFoundError, ValueError) as exc: + _baseline_load_error = str(exc) + logger.error( + "gate baseline could not be loaded from %s: %s -- every agent will read as " + "unknown until this is fixed and the process restarts", + baseline_path, + exc, + ) + else: + logger.warning( + "DUSK_GATE_BASELINE_PATH not set; gate has no baseline, every agent is unknown" + ) + return gate_engine + + +def _get_gate_engine() -> ActionGate: + # Live traffic never updates the trusted baseline; doing so would permit + # gradual baseline poisoning. + global _gate_engine + if _gate_engine is None: + with _gate_lock: + if _gate_engine is None: + _gate_engine = _load_gate_engine() + return _gate_engine + + +def reset_gate_engine() -> None: + """Clear the cached gate engine so the next request reloads it. Test-only hook.""" + global _gate_engine + with _gate_lock: + _gate_engine = None + + +#: Capped decision history with embeddings computed at record time. +_DECISION_HISTORY_CAP = 200 +#: Per-agent sub-cap so one noisy agent can't evict the whole fleet's candidates. +_DECISION_HISTORY_PER_AGENT_CAP = 40 +_decision_history: list[tuple[TraceDecision, list[float]]] = [] +_decision_history_lock = threading.Lock() + + +def reset_decision_history() -> None: + """Clear recorded decisions used for similarity lookups. Test-only hook.""" + with _decision_history_lock: + _decision_history.clear() + + +def _find_similar_decisions(agent_id: str, action_text: str) -> list[str]: + from dusk.trace.vector import find_similar_cached + + with _decision_history_lock: + history_snapshot = list(_decision_history) + similar = find_similar_cached(action_text, agent_id, history_snapshot, top_k=3) + return [s.id for s in similar] + + +def _record_decision( + decision_id: str, + agent_id: str, + action_text: str, + score: float, + reasons: list[str], + similar_decision_ids: list[str], + verdict: str, +) -> None: + from dusk.trace.models import TraceDecision + from dusk.trace.vector import embed_text + + decision = TraceDecision( + id=decision_id, + agent_id=agent_id, + action=action_text, + score=round(score * 100), + reasoning=reasons[0] if reasons else "", + risk_flags=reasons, + similar_decision_ids=similar_decision_ids, + verdict=verdict, + ) + # Candidate and stored embeddings must use the same text shape. + vec = embed_text(f"{agent_id} {action_text}") + with _decision_history_lock: + agent_indices = [i for i, (d, _v) in enumerate(_decision_history) if d.agent_id == agent_id] + if len(agent_indices) >= _DECISION_HISTORY_PER_AGENT_CAP: + del _decision_history[agent_indices[0]] + _decision_history.append((decision, vec)) + if len(_decision_history) > _DECISION_HISTORY_CAP: + del _decision_history[: len(_decision_history) - _DECISION_HISTORY_CAP] + + +@app.route("/v1/gate", methods=["POST"]) +def evaluate_gate_action() -> object: + """Evaluate a proposed agent action against the learned baseline. + + Contract: contracts/gate.openapi.yaml. + """ + from dusk.actions.event import AgentAction + + raw = request.get_json(force=True, silent=True) + if not isinstance(raw, dict): + return jsonify({"error": "request body must be a JSON object"}), 400 + + try: + action = AgentAction.from_dict(raw) + except ValueError as exc: + return jsonify({"error": str(exc)}), 400 + + verdict = _get_gate_engine().evaluate(action) + analysis = verdict.analysis + action_text = f"{action.action_type} {action.target}" + trace_id = verdict.trace_id + similar_decision_ids = _find_similar_decisions(action.agent_id, action_text) + + response: dict[str, object] = { + "trace_id": trace_id, + "verdict": verdict.verdict, + "score": round(analysis.score, 4), + "blast": analysis.blast_radius, + "mitre_attack": [analysis.mitre_attack] if analysis.mitre_attack else [], + "mitre_atlas": [analysis.mitre_atlas] if analysis.mitre_atlas else [], + "reasons": analysis.reasons, + "predicted_next": analysis.predicted_next, + "similar_decision_ids": similar_decision_ids, + } + logger.info( + "gate verdict trace_id=%s agent=%s verdict=%s score=%.2f", + response["trace_id"], + action.agent_id, + verdict.verdict, + analysis.score, + ) + _record_decision( + trace_id, + action.agent_id, + action_text, + analysis.score, + analysis.reasons, + similar_decision_ids, + verdict.verdict, + ) + + from dusk.trace.n8n_client import fire_alert, fire_decision, fire_report + + webhook_payload = { + **response, + "agent_id": action.agent_id, + "action_type": action.action_type, + "target": action.target, + } + fire_decision(webhook_payload) + fire_report(webhook_payload) + if verdict.refused: + fire_alert(webhook_payload) + + return jsonify(response), 200 + + +@app.route("/health") +def health() -> object: + gate_engine = _get_gate_engine() # forces the baseline load attempt so it's reflected below + if _baseline_load_error is not None: + # A non-200 response propagates the failure to the container health check. + return jsonify({"status": "degraded", "baseline_error": _baseline_load_error}), 503 + offense_memory = gate_engine.offense_memory + persist_error = offense_memory.last_persist_error if offense_memory is not None else None + if persist_error is not None: + # The gate remains available, but repeat-offense durability is degraded. + return jsonify({"status": "degraded", "offense_memory_error": persist_error}), 503 + return jsonify({"status": "ok"}) + + +def run() -> None: + port = int(os.getenv("FLASK_PORT", "5000")) + host = os.getenv("FLASK_HOST", "127.0.0.1") + app.run(host=host, port=port) + + +if __name__ == "__main__": + run() diff --git a/examples/agent-action-monitor/src/dusk/config.py b/examples/agent-action-monitor/src/dusk/config.py new file mode 100644 index 000000000..398f76caf --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/config.py @@ -0,0 +1,286 @@ +"""Configuration system for Dusk. + +Configuration is resolved from three layers, lowest to highest precedence: + +1. Dataclass defaults defined on :class:`Config`. +2. A ``dusk.yaml`` file in the current working directory, if present. +3. Environment variables prefixed ``DUSK_`` (e.g. ``DUSK_SWEEP_THRESHOLD``). + +The active configuration is a process-wide singleton: :func:`get_config` +loads it once and caches it. Detections and the engine receive a +:class:`Config` instance rather than reading any hardcoded values, so every +threshold has a single, overridable source of truth. +""" + +from __future__ import annotations + +import dataclasses +import logging +import os +from dataclasses import dataclass +from typing import Any, get_type_hints + +import yaml + +logger = logging.getLogger("dusk.config") + +#: Prefix for environment-variable overrides. +ENV_PREFIX = "DUSK_" +#: Default config file looked up in the current working directory. +CONFIG_FILENAME = "dusk.yaml" + + +class ConfigError(Exception): + """Raised when configuration values are invalid or cannot be loaded.""" + + +@dataclass(frozen=True) +class Config: + """Immutable Dusk configuration. + + Attributes: + sweep_threshold: Unique destinations within the sweep window above + which a source is considered to be sweeping. + sweep_window_seconds: Length of the sweep sliding window, in seconds. + sweep_timing_std_threshold: Inter-packet interval standard deviation + (seconds) below which timing is judged machine-regular. + boundary_port_threshold: Unique destination ports for one + ``(src, dst)`` pair above which a port scan is flagged. + boundary_window_seconds: Length of the boundary sliding window, in + seconds. + gate_block_threshold: Anomaly score (0..1) at or above which the agent + action gate refuses an action (WOULD-BLOCK in watch mode, BLOCK in + enforce mode). + alert_log_path: Path the alert responder appends JSON entries to. + log_level: Default logging level name (e.g. ``"WARNING"``). + enforce: When ``True``, the ``/v1/gate`` HTTP service renders BLOCK + for refused actions instead of WOULD-BLOCK. Watch mode (``False``) + is the default. + sie_endpoint: Base URL of the self-hosted SIE container the gate + calls for encode/score/extract. No API key is needed for a + self-hosted instance. + sie_encode_model: Catalog model id SIE uses for the encode primitive. + sie_score_model: Catalog model id SIE uses for the score (rerank) + primitive. + sie_extract_model: Catalog model id SIE uses for the extract + primitive. + sie_timeout_ms: Timeout in milliseconds for a single SIE call. + n8n_alert_url: n8n webhook fired only when a gate verdict is refused. + n8n_report_url: n8n webhook fired on every gate verdict. + n8n_decision_url: n8n webhook fired on every gate verdict, kept on + its own URL so automation and reporting can be routed separately. + n8n_max_workers: Size of the bounded thread pool webhook firing uses, + so a sustained burst of refused verdicts can't spawn unbounded + OS threads. + n8n_max_queued: Webhook sends allowed to be queued or in flight at + once; past this, new ones are dropped and logged instead of + growing the backlog without limit. + offense_memory_path: File the gate persists its repeat-offense + memory to, so it survives a service restart. Empty (the + default) disables persistence -- the signal still scores within + one process lifetime, from an in-memory-only store, but nothing + is written to disk. Deliberately opt-in, matching the + ``n8n_*_url`` fields: an operator running the service outside + the Docker Compose stack (which sets this explicitly) shouldn't + get a file silently dropped into whatever the process's current + working directory happens to be. + repeat_offense_max_contribution: Upper bound on how much the + repeat-offense signal alone can add to an anomaly score, so it + is one input among several rather than something that alone can + force a BLOCK on an otherwise benign action. + repeat_offense_half_life_days: Time, in days, after which a past + offense's contribution to the repeat-offense signal is halved. + Models the intuition that an agent blocked once six months ago + is not the same risk as one blocked five minutes ago. + """ + + sweep_threshold: int = 15 + sweep_window_seconds: float = 10.0 + sweep_timing_std_threshold: float = 0.05 + boundary_port_threshold: int = 10 + boundary_window_seconds: float = 30.0 + gate_block_threshold: float = 0.6 + alert_log_path: str = "dusk-alerts.json" + log_level: str = "WARNING" + enforce: bool = False + sie_endpoint: str = "http://sie:8080" + sie_encode_model: str = "BAAI/bge-m3" + sie_score_model: str = "BAAI/bge-reranker-v2-m3" + sie_extract_model: str = "urchade/gliner_multi-v2.1" + sie_timeout_ms: int = 10000 + n8n_alert_url: str = "" + n8n_report_url: str = "" + n8n_decision_url: str = "" + n8n_max_workers: int = 8 + n8n_max_queued: int = 200 + offense_memory_path: str = "" + repeat_offense_max_contribution: float = 0.3 + repeat_offense_half_life_days: float = 30.0 + + def __post_init__(self) -> None: + """Validate values after construction. + + Raises: + ConfigError: If any numeric threshold is non-positive, the log + level is not a recognised name, or a configured URL is + non-empty but missing an http(s) scheme. + """ + positive_numeric = ( + "sweep_threshold", + "sweep_window_seconds", + "sweep_timing_std_threshold", + "boundary_port_threshold", + "boundary_window_seconds", + "gate_block_threshold", + "sie_timeout_ms", + "n8n_max_workers", + "n8n_max_queued", + "repeat_offense_max_contribution", + "repeat_offense_half_life_days", + ) + for name in positive_numeric: + value = getattr(self, name) + if value <= 0: + raise ConfigError(f"Config value '{name}' must be > 0, got {value!r}") + + if not isinstance(logging.getLevelName(self.log_level), int): + raise ConfigError( + f"Config value 'log_level' is not a valid level name: {self.log_level!r}" + ) + + url_fields = ("sie_endpoint", "n8n_alert_url", "n8n_report_url", "n8n_decision_url") + for name in url_fields: + value = getattr(self, name) + if value and not value.startswith(("http://", "https://")): + raise ConfigError( + f"Config value '{name}' must be empty or start with http:// or " + f"https://, got {value!r}" + ) + + def to_dict(self) -> dict[str, Any]: + """Return a plain-dict view of the configuration.""" + return dataclasses.asdict(self) + + +def _coerce(field_type: Any, raw: Any, source: str) -> Any: # noqa: ANN401 + """Coerce ``raw`` to ``field_type``, raising :class:`ConfigError` on failure. + + Args: + field_type: The dataclass field's declared type (``int``, ``float`` + or ``str``). + raw: The value read from YAML or the environment. + source: Human-readable origin used in error messages. + + Returns: + ``raw`` converted to ``field_type``. + + Raises: + ConfigError: If ``raw`` cannot be coerced to ``field_type``. + """ + try: + if field_type is bool: + return str(raw).strip().lower() in {"1", "true", "yes", "on"} + if field_type is int: + return int(raw) + if field_type is float: + return float(raw) + return str(raw) + except (TypeError, ValueError) as exc: + raise ConfigError(f"Invalid value for {source}: {raw!r} ({exc})") from exc + + +def _load_yaml_overrides(path: str) -> dict[str, Any]: + """Read overrides from a ``dusk.yaml`` file, returning {} if absent. + + Raises: + ConfigError: If the file exists but cannot be parsed or is not a + mapping. + """ + if not os.path.exists(path): + return {} + + try: + with open(path, encoding="utf-8") as handle: + loaded = yaml.safe_load(handle) + except (OSError, yaml.YAMLError) as exc: + raise ConfigError(f"Could not read config file '{path}': {exc}") from exc + + if loaded is None: + return {} + if not isinstance(loaded, dict): + raise ConfigError( + f"Config file '{path}' must contain a mapping, got {type(loaded).__name__}" + ) + + logger.debug("Loaded %d override(s) from %s", len(loaded), path) + return loaded + + +def _env_overrides(field_names: set[str]) -> dict[str, str]: + """Collect ``DUSK_*`` environment overrides for known fields.""" + overrides: dict[str, str] = {} + for env_name, value in os.environ.items(): + if not env_name.startswith(ENV_PREFIX): + continue + field_name = env_name[len(ENV_PREFIX) :].lower() + if field_name in field_names: + overrides[field_name] = value + logger.debug("Applied env override %s", env_name) + return overrides + + +def load_config(config_path: str = CONFIG_FILENAME) -> Config: + """Build a :class:`Config` from defaults, ``dusk.yaml`` and env vars. + + Args: + config_path: Path to the YAML config file. Defaults to ``dusk.yaml`` + in the current working directory. + + Returns: + A validated :class:`Config` instance. + + Raises: + ConfigError: If a config file is malformed or a value is invalid. + """ + # get_type_hints resolves real type objects even though this module uses + # `from __future__ import annotations` (which would otherwise yield strings). + type_by_field = get_type_hints(Config) + field_names = set(type_by_field) + + merged: dict[str, Any] = {} + + for field_name, raw in _load_yaml_overrides(config_path).items(): + if field_name not in field_names: + logger.warning("Ignoring unknown config key in %s: %s", config_path, field_name) + continue + merged[field_name] = _coerce(type_by_field[field_name], raw, f"{config_path}:{field_name}") + + for field_name, raw in _env_overrides(field_names).items(): + merged[field_name] = _coerce( + type_by_field[field_name], raw, f"{ENV_PREFIX}{field_name.upper()}" + ) + + return Config(**merged) + + +_active_config: Config | None = None + + +def get_config() -> Config: + """Return the process-wide singleton config, loading it on first use.""" + global _active_config + if _active_config is None: + _active_config = load_config() + return _active_config + + +def set_config(config: Config) -> None: + """Replace the active singleton config (primarily for tests and the CLI).""" + global _active_config + _active_config = config + + +def reset_config() -> None: + """Clear the cached singleton so the next :func:`get_config` reloads it.""" + global _active_config + _active_config = None diff --git a/examples/agent-action-monitor/src/dusk/trace/__init__.py b/examples/agent-action-monitor/src/dusk/trace/__init__.py new file mode 100644 index 000000000..109a41cf4 --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/trace/__init__.py @@ -0,0 +1 @@ +"""TRACE -- audit trail layer for DUSK.""" diff --git a/examples/agent-action-monitor/src/dusk/trace/models.py b/examples/agent-action-monitor/src/dusk/trace/models.py new file mode 100644 index 000000000..b14083db6 --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/trace/models.py @@ -0,0 +1,59 @@ +"""TraceDecision -- the canonical audit record for one agent action.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from uuid import uuid4 + + +@dataclass +class TraceDecision: + """One recorded agent decision with full audit trail. + + Fields agent_id, action, score, and reasoning are required. + All others default so that api.py can construct quickly and augment later. + """ + + agent_id: str + action: str + score: int + reasoning: str + id: str = field(default_factory=lambda: uuid4().hex[:8]) + timestamp: float = field(default_factory=time.time) + risk_flags: list[str] = field(default_factory=list) + similar_decision_ids: list[str] = field(default_factory=list) + #: The gate's actual verdict (ALLOW / WOULD-BLOCK / BLOCK) for this decision, empty when + #: unknown (e.g. a decision recorded before this field existed). + verdict: str = "" + + @property + def risk_level(self) -> str: + if self.score >= 70: + return "high" + if self.score >= 40: + return "medium" + return "low" + + def to_dict(self) -> dict[str, object]: + return { + "id": self.id, + "agent_id": self.agent_id, + "action": self.action, + "score": self.score, + "reasoning": self.reasoning, + "risk_flags": self.risk_flags, + "timestamp": self.timestamp, + "verdict": self.verdict, + "output": { + "score": self.score, + "reasoning": self.reasoning, + "confidence": round(self.score / 100, 2), + "risk_flags": self.risk_flags, + }, + "trace": { + "status": "recorded", + "risk_level": self.risk_level, + "similar_decisions": self.similar_decision_ids, + }, + } diff --git a/examples/agent-action-monitor/src/dusk/trace/n8n_client.py b/examples/agent-action-monitor/src/dusk/trace/n8n_client.py new file mode 100644 index 000000000..bccdd5a4f --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/trace/n8n_client.py @@ -0,0 +1,96 @@ +"""Non-blocking n8n decision, report, and alert webhooks.""" + +from __future__ import annotations + +import json +import logging +import threading +import urllib.error +import urllib.request +from concurrent.futures import ThreadPoolExecutor + +from dusk.config import Config, get_config + +logger = logging.getLogger(__name__) + +#: Pool size is fixed from process configuration on first use. +_executor: ThreadPoolExecutor | None = None + +#: Queued-or-in-flight send count; bounds the pool's own unbounded task queue. +_pending_lock = threading.Lock() +_pending_count = 0 + + +def _get_executor() -> ThreadPoolExecutor: + global _executor + if _executor is None: + _executor = ThreadPoolExecutor( + max_workers=get_config().n8n_max_workers, thread_name_prefix="n8n-webhook" + ) + return _executor + + +def _submit(url: str, label: str, payload: dict[str, object], config: Config) -> None: + if not url: + return + global _pending_count + with _pending_lock: + if _pending_count >= config.n8n_max_queued: + logger.warning( + "n8n webhook (%s) dropped: %d already queued or in flight (n8n_max_queued=%d)", + label, + _pending_count, + config.n8n_max_queued, + ) + return + _pending_count += 1 + _get_executor().submit(_send_and_release, url, label, payload) + + +def _send_and_release(url: str, label: str, payload: dict[str, object]) -> None: + global _pending_count + try: + _send(url, label, payload) + finally: + with _pending_lock: + _pending_count -= 1 + + +def fire_decision(payload: dict[str, object], config: Config | None = None) -> None: + """Fire the decision webhook on the bounded pool -- never blocks the caller.""" + cfg = config or get_config() + _submit(cfg.n8n_decision_url, "decision", payload, cfg) + + +def fire_report(payload: dict[str, object], config: Config | None = None) -> None: + """Fire the report webhook on the bounded pool -- never blocks the caller.""" + cfg = config or get_config() + _submit(cfg.n8n_report_url, "report", payload, cfg) + + +def fire_alert(payload: dict[str, object], config: Config | None = None) -> None: + """Fire the alert webhook on the bounded pool -- never blocks the caller.""" + cfg = config or get_config() + _submit(cfg.n8n_alert_url, "alert", payload, cfg) + + +def _send(url: str, label: str, payload: dict[str, object]) -> None: + if not url: + return + if not url.startswith(("https://", "http://")): + logger.warning("n8n webhook (%s) has unsupported scheme, skipping", label) + return + try: + body = json.dumps(payload).encode() + req = urllib.request.Request( # noqa: S310 + url, + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=5) as resp: # noqa: S310 # nosec B310 + logger.info("n8n webhook (%s) fired, status=%s", label, resp.status) + except urllib.error.URLError as exc: + logger.warning("n8n webhook (%s) failed: %s", label, exc) + except Exception as exc: # noqa: BLE001 + logger.warning("n8n webhook (%s) error: %s", label, exc) diff --git a/examples/agent-action-monitor/src/dusk/trace/vector.py b/examples/agent-action-monitor/src/dusk/trace/vector.py new file mode 100644 index 000000000..de86fd1d4 --- /dev/null +++ b/examples/agent-action-monitor/src/dusk/trace/vector.py @@ -0,0 +1,304 @@ +"""SIE-backed similarity with deterministic offline fallbacks.""" + +from __future__ import annotations + +import hashlib +import logging +import math +import os +from dataclasses import dataclass +from typing import Any + +from dusk.config import Config, get_config +from dusk.trace.models import TraceDecision + +logger = logging.getLogger(__name__) + +#: Default zero-shot labels for pulling privileged terms out of an action. +DEFAULT_EXTRACT_LABELS = ["role", "privilege", "resource", "segment", "port"] + +#: Per-call provisioning bound; one gate request can make several sequential SIE calls. +_PROVISION_TIMEOUT_S = 1.5 + + +def _sigmoid(x: float) -> float: + """Bound a raw cross-encoder logit into [0, 1]; monotonic, not a calibrated probability.""" + try: + return 1.0 / (1.0 + math.exp(-x)) + except OverflowError: + return 0.0 if x < 0 else 1.0 + + +@dataclass +class ExtractedTerm: + """One entity GLiNER pulled out of an action's text, with its confidence.""" + + text: str + label: str + score: float + + +@dataclass +class SimilarDecision: + id: str + agent_id: str + action: str + similarity: float + verdict: str + score: int + + +#: Conservative label for records created before verdict persistence existed. +_UNKNOWN_VERDICT_FALLBACK = "WOULD-BLOCK" + + +def _sie_client(config: Config) -> Any | None: # noqa: ANN401 + """Return a constructed SIEClient if the sie-sdk package is installed, else None.""" + try: + from sie_sdk import SIEClient # type: ignore[import-not-found] + except ImportError: + return None + try: + api_key = os.getenv("SIE_API_KEY") or None + return SIEClient( + config.sie_endpoint.rstrip("/"), + api_key=api_key, + timeout_s=config.sie_timeout_ms / 1000, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("SIE client construction failed: %s", exc) + return None + + +def sie_encode(text: str, config: Config | None = None) -> list[float] | None: + """encode: text -> dense vector via SIE. Returns None to trigger the n-gram fallback.""" + cfg = config or get_config() + client = _sie_client(cfg) + if client is None: + return None + try: + from sie_sdk.types import Item # type: ignore[import-not-found] + + # Cold or saturated models must not stall the inline gate. + result = client.encode( + cfg.sie_encode_model, + Item(text=text), + wait_for_capacity=False, + provision_timeout_s=_PROVISION_TIMEOUT_S, + ) + dense = result["dense"] if isinstance(result, dict) else getattr(result, "dense", None) + return [float(v) for v in dense] if dense is not None else None + except Exception as exc: # noqa: BLE001 + logger.warning("SIE encode failed: %s", exc) + return None + + +def sie_score( + query: str, candidates: list[str], config: Config | None = None +) -> list[float] | None: + """score: rerank candidates against query via SIE's cross-encoder. + + Returns one [0, 1]-bounded score per candidate (see :func:`_sigmoid`) in + the same order as ``candidates`` (never reordered), or None when SIE is + unavailable or there are no candidates to score. + """ + if not candidates: + return None + cfg = config or get_config() + client = _sie_client(cfg) + if client is None: + return None + try: + from sie_sdk.types import Item # type: ignore[import-not-found] + + query_item = Item(text=query) + candidate_items = [Item(text=text, id=str(i)) for i, text in enumerate(candidates)] + result = client.score( + cfg.sie_score_model, + query_item, + candidate_items, + wait_for_capacity=False, + provision_timeout_s=_PROVISION_TIMEOUT_S, + ) + entries = result["scores"] if isinstance(result, dict) else getattr(result, "scores", None) + if not entries: + return None + score_by_id = {str(e["item_id"]): _sigmoid(float(e["score"])) for e in entries} + return [score_by_id.get(str(i), 0.0) for i in range(len(candidates))] + except Exception as exc: # noqa: BLE001 + logger.warning("SIE score failed: %s", exc) + return None + + +def sie_extract( + text: str, labels: list[str] | None = None, config: Config | None = None +) -> list[ExtractedTerm]: + """extract: pull entities / privileged terms from text via SIE's GLiNER model. + + Returns an empty list when SIE is unavailable, never raises. Each result + keeps its label and confidence score so callers can weight by how + confident this zero-shot extraction actually was, rather than treating + every hit as equally privileged. + """ + cfg = config or get_config() + client = _sie_client(cfg) + if client is None: + return [] + try: + from sie_sdk.types import Item # type: ignore[import-not-found] + + item_labels = labels or DEFAULT_EXTRACT_LABELS + result = client.extract( + cfg.sie_extract_model, + Item(text=text), + labels=item_labels, + wait_for_capacity=False, + provision_timeout_s=_PROVISION_TIMEOUT_S, + ) + entities = ( + result["entities"] if isinstance(result, dict) else getattr(result, "entities", None) + ) + if not entities: + return [] + extracted: list[ExtractedTerm] = [] + for e in entities: + if e is None: + continue + if isinstance(e, dict): + extracted.append( + ExtractedTerm( + text=str(e["text"]), + label=str(e.get("label", "")), + score=float(e.get("score", 0.0)), + ) + ) + else: + extracted.append( + ExtractedTerm( + text=str(e.text), + label=str(getattr(e, "label", "")), + score=float(getattr(e, "score", 0.0)), + ) + ) + return extracted + except Exception as exc: # noqa: BLE001 + logger.warning("SIE extract failed: %s", exc) + return [] + + +def _stable_hash(token: str) -> int: + """Return a process-stable token hash for persisted fallback embeddings.""" + return int(hashlib.sha256(token.encode("utf-8")).hexdigest(), 16) + + +def _ngram_fallback(text: str, dims: int = 64) -> list[float]: + vec = [0.0] * dims + for token in text.lower().split(): + vec[_stable_hash(token) % dims] += 1.0 + norm = math.sqrt(sum(v * v for v in vec)) or 1.0 + return [v / norm for v in vec] + + +def _cosine(a: list[float], b: list[float]) -> float: + if len(a) != len(b) or not a: + return 0.0 + dot = sum(x * y for x, y in zip(a, b, strict=True)) + mag = math.sqrt(sum(x * x for x in a)) * math.sqrt(sum(y * y for y in b)) + return dot / mag if mag else 0.0 + + +def embed_text(text: str, config: Config | None = None) -> list[float]: + """Return a dense vector for text via SIE encode, or the n-gram fallback.""" + return sie_encode(text, config) or _ngram_fallback(text) + + +def _rank_candidates( + query: str, + scored: list[tuple[float, TraceDecision]], + top_k: int, +) -> list[SimilarDecision]: + """Shared top_k selection + optional SIE rerank pass, given pre-scored candidates.""" + scored.sort(key=lambda x: x[0], reverse=True) + top = scored[:top_k] + + # Prefer the cross-encoder ordering when SIE can rerank the shortlist. + rerank = sie_score(query, [f"{d.agent_id} {d.action}" for _, d in top]) + if rerank and len(rerank) == len(top): + reranked = sorted(zip(rerank, top, strict=True), key=lambda x: x[0], reverse=True) + top = [t for _, t in reranked] + + return [ + SimilarDecision( + id=d.id, + agent_id=d.agent_id, + action=d.action, + similarity=round(sim, 3), + verdict=d.verdict or _UNKNOWN_VERDICT_FALLBACK, + score=d.score, + ) + for sim, d in top + ] + + +def find_similar( + target_action: str, + target_agent: str, + past_decisions: list[TraceDecision], + top_k: int = 3, +) -> list[SimilarDecision]: + """Return the top_k most similar past decisions to a new one. + + Re-embeds every past decision on every call -- fine for an occasional, + small lookup, but O(len(past_decisions)) SIE encode calls. A live + request path with growing history should use :func:`find_similar_cached` + instead, which embeds each decision once at record time. + + Returns an empty list when fewer than 2 past decisions exist. + Never raises -- worst case is an empty list. + """ + if len(past_decisions) < 2: + return [] + + query = f"{target_agent} {target_action}" + query_vec = embed_text(query) + + scored: list[tuple[float, TraceDecision]] = [] + for d in past_decisions: + candidate = f"{d.agent_id} {d.action} {d.reasoning}" + candidate_vec = embed_text(candidate) + sim = _cosine(query_vec, candidate_vec) + if sim > 0.3: + scored.append((sim, d)) + + return _rank_candidates(query, scored, top_k) + + +def find_similar_cached( + target_action: str, + target_agent: str, + past_decisions: list[tuple[TraceDecision, list[float]]], + top_k: int = 3, +) -> list[SimilarDecision]: + """Like find_similar, but each past decision already carries its embedding. + + Costs exactly one SIE encode call (the query) plus an optional rerank on + the shortlist, independent of how much history exists -- unlike + find_similar, which re-embeds the entire history on every call. Intended + for a live request path where each decision's vector is computed once, + when it's recorded, and reused for every future lookup. + + Returns an empty list when fewer than 2 past decisions exist. + """ + if len(past_decisions) < 2: + return [] + + query = f"{target_agent} {target_action}" + query_vec = embed_text(query) + + scored: list[tuple[float, TraceDecision]] = [] + for decision, vec in past_decisions: + sim = _cosine(query_vec, vec) + if sim > 0.3: + scored.append((sim, decision)) + + return _rank_candidates(query, scored, top_k) diff --git a/examples/agent-action-monitor/tests/conftest.py b/examples/agent-action-monitor/tests/conftest.py new file mode 100644 index 000000000..806d89fc9 --- /dev/null +++ b/examples/agent-action-monitor/tests/conftest.py @@ -0,0 +1,21 @@ +"""Shared pytest fixtures for the Dusk test suite.""" + +from __future__ import annotations + +from collections.abc import Iterator + +import pytest + +from dusk.config import reset_config + + +@pytest.fixture(autouse=True) +def _isolate_config() -> Iterator[None]: + """Reset the config singleton before and after every test. + + Ensures one test's ``set_config`` / environment overrides never leak into + another, keeping detection thresholds deterministic. + """ + reset_config() + yield + reset_config() diff --git a/examples/agent-action-monitor/tests/integration/__init__.py b/examples/agent-action-monitor/tests/integration/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/examples/agent-action-monitor/tests/integration/test_gate_api.py b/examples/agent-action-monitor/tests/integration/test_gate_api.py new file mode 100644 index 000000000..897b49eac --- /dev/null +++ b/examples/agent-action-monitor/tests/integration/test_gate_api.py @@ -0,0 +1,322 @@ +"""Tests for the /v1/gate HTTP endpoint (contracts/gate.openapi.yaml).""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest + +from dusk import api +from dusk.config import reset_config + +FIXTURES = Path(__file__).resolve().parent.parent / "fixtures" +BASELINE_PATH = str(FIXTURES / "actions_normal.json") + +LAB_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "lab", "actions") +sys.path.insert(0, os.path.abspath(LAB_DIR)) + +import generate_actions # noqa: E402 + +if not Path(BASELINE_PATH).exists(): + generate_actions.generate(str(FIXTURES)) + +CONTRACT_FIELDS = { + "trace_id", + "verdict", + "score", + "blast", + "mitre_attack", + "mitre_atlas", + "reasons", + "predicted_next", + "similar_decision_ids", +} + + +@pytest.fixture(autouse=True) +def _reset_gate(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + monkeypatch.setenv("DUSK_GATE_BASELINE_PATH", BASELINE_PATH) + monkeypatch.setenv("DUSK_GATE_BASELINE_SOURCE", "generic") + monkeypatch.delenv("DUSK_ENFORCE", raising=False) + # Isolated per test by default so a refusal in one test never writes + # into the repo's working directory; test_offense_memory_persists_* + # overrides this explicitly to exercise the real default-path behaviour. + monkeypatch.setenv("DUSK_OFFENSE_MEMORY_PATH", str(tmp_path / "offense-memory.json")) + reset_config() + api.reset_gate_engine() + api.reset_decision_history() + yield + reset_config() + api.reset_gate_engine() + api.reset_decision_history() + + +@pytest.fixture +def client(): + api.app.config["TESTING"] = True + with api.app.test_client() as c: + yield c + + +def _action_payload( + agent_id: str = "netops-agent", target: str = "fw-corp-https", **after: object +) -> dict[str, object]: + return { + "agent_id": agent_id, + "timestamp": "2023-11-14T22:20:00+00:00", + "action_type": "firewall_rule_change", + "target": target, + "change": {"before": None, "after": dict(after) if after else None}, + "source": "generic", + "raw_ref": "evt-test-1", + } + + +def test_health(client) -> None: + r = client.get("/health") + assert r.status_code == 200 + assert r.get_json()["status"] == "ok" + + +def test_gate_returns_contract_shaped_verdict(client) -> None: + r = client.post("/v1/gate", json=_action_payload(port=443)) + assert r.status_code == 200 + data = r.get_json() + assert set(data) == CONTRACT_FIELDS + assert data["verdict"] in {"ALLOW", "WOULD-BLOCK", "BLOCK"} + assert isinstance(data["mitre_attack"], list) + assert isinstance(data["mitre_atlas"], list) + assert isinstance(data["reasons"], list) + assert isinstance(data["similar_decision_ids"], list) + assert 0.0 <= data["score"] <= 1.0 + assert data["blast"] in {"low", "medium", "high"} + + +def test_gate_allows_known_agent_pattern(client) -> None: + r = client.post("/v1/gate", json=_action_payload(port=443)) + assert r.get_json()["verdict"] == "ALLOW" + + +def test_gate_flags_unknown_agent_touching_sensitive_target(client) -> None: + r = client.post( + "/v1/gate", json=_action_payload(agent_id="ghost-agent", target="fw-restricted") + ) + assert r.get_json()["verdict"] in {"WOULD-BLOCK", "BLOCK"} + + +def test_gate_rejects_invalid_action(client) -> None: + r = client.post("/v1/gate", json={"agent_id": "netops-agent"}) + assert r.status_code == 400 + assert "error" in r.get_json() + + +def test_gate_rejects_non_object_body(client) -> None: + r = client.post("/v1/gate", data="not json", content_type="application/json") + assert r.status_code == 400 + + +def test_gate_rejects_oversized_body(client) -> None: + oversized = "x" * (2 * 1024 * 1024) + r = client.post("/v1/gate", data=oversized, content_type="application/json") + assert r.status_code == 413 + + +def test_gate_without_baseline_defaults_to_unknown_agent(client, monkeypatch) -> None: + monkeypatch.delenv("DUSK_GATE_BASELINE_PATH", raising=False) + api.reset_gate_engine() + r = client.post("/v1/gate", json=_action_payload()) + assert r.status_code == 200 + assert any("no established baseline" in reason for reason in r.get_json()["reasons"]) + + +def test_health_reports_ok_when_baseline_not_configured(client, monkeypatch) -> None: + monkeypatch.delenv("DUSK_GATE_BASELINE_PATH", raising=False) + api.reset_gate_engine() + r = client.get("/health") + assert r.status_code == 200 + assert r.get_json()["status"] == "ok" + + +def test_health_reports_degraded_when_baseline_path_is_broken(client, monkeypatch) -> None: + monkeypatch.setenv("DUSK_GATE_BASELINE_PATH", "/does/not/exist.json") + api.reset_gate_engine() + r = client.get("/health") + assert r.status_code == 503 + body = r.get_json() + assert body["status"] == "degraded" + assert "baseline_error" in body + + +def test_gate_still_serves_requests_when_baseline_is_broken(client, monkeypatch) -> None: + """/v1/gate degrades (every agent unknown) rather than refusing outright.""" + monkeypatch.setenv("DUSK_GATE_BASELINE_PATH", "/does/not/exist.json") + api.reset_gate_engine() + r = client.post("/v1/gate", json=_action_payload()) + assert r.status_code == 200 + assert any("no established baseline" in reason for reason in r.get_json()["reasons"]) + + +def test_health_reports_degraded_when_offense_memory_persistence_fails( + client, tmp_path, monkeypatch +) -> None: + """A silently failing repeat-offense write must be visible to monitoring, not just logs.""" + blocker = tmp_path / "not-a-directory" + blocker.write_text("blocking file", encoding="utf-8") + bad_storage = blocker / "offense-memory.json" + monkeypatch.setenv("DUSK_OFFENSE_MEMORY_PATH", str(bad_storage)) + reset_config() + api.reset_gate_engine() + + payload = _action_payload(agent_id="ghost-agent", target="fw-restricted") + r = client.post("/v1/gate", json=payload) + assert r.get_json()["verdict"] in {"WOULD-BLOCK", "BLOCK"} + api._get_gate_engine().offense_memory.flush() + + health = client.get("/health") + assert health.status_code == 503 + body = health.get_json() + assert body["status"] == "degraded" + assert "offense_memory_error" in body + + +def test_gate_enforce_mode_via_config_blocks_instead_of_would_block(client, monkeypatch) -> None: + monkeypatch.setenv("DUSK_ENFORCE", "true") + reset_config() + api.reset_gate_engine() + r = client.post( + "/v1/gate", json=_action_payload(agent_id="ghost-agent", target="fw-restricted") + ) + assert r.get_json()["verdict"] == "BLOCK" + + +def test_gate_allow_fires_decision_and_report_but_not_alert(client) -> None: + with ( + patch("dusk.trace.n8n_client.fire_decision") as mock_decision, + patch("dusk.trace.n8n_client.fire_report") as mock_report, + patch("dusk.trace.n8n_client.fire_alert") as mock_alert, + ): + r = client.post("/v1/gate", json=_action_payload(port=443)) + + assert r.get_json()["verdict"] == "ALLOW" + mock_decision.assert_called_once() + mock_report.assert_called_once() + mock_alert.assert_not_called() + + +def test_gate_refusal_fires_all_three_webhooks(client) -> None: + with ( + patch("dusk.trace.n8n_client.fire_decision") as mock_decision, + patch("dusk.trace.n8n_client.fire_report") as mock_report, + patch("dusk.trace.n8n_client.fire_alert") as mock_alert, + ): + r = client.post( + "/v1/gate", json=_action_payload(agent_id="ghost-agent", target="fw-restricted") + ) + + assert r.get_json()["verdict"] in {"WOULD-BLOCK", "BLOCK"} + mock_decision.assert_called_once() + mock_report.assert_called_once() + mock_alert.assert_called_once() + + +def test_gate_webhook_payload_includes_action_context(client) -> None: + with patch("dusk.trace.n8n_client.fire_decision") as mock_decision: + client.post("/v1/gate", json=_action_payload(port=443)) + + payload = mock_decision.call_args[0][0] + assert payload["agent_id"] == "netops-agent" + assert payload["action_type"] == "firewall_rule_change" + assert payload["target"] == "fw-corp-https" + assert set(CONTRACT_FIELDS) <= set(payload) + + +def test_recorded_decision_carries_real_risk_flags_and_similar_ids(client) -> None: + """TraceDecision.risk_flags and .similar_decision_ids must reflect what the + response actually computed, not stay at their dataclass defaults -- both + were previously always empty on the stored object even when the response + carried real values (similar_decision_ids) or real reasons existed + (risk_flags). find_similar_cached needs at least 2 prior decisions before + it returns anything, so this fires three near-identical actions and + checks the third.""" + client.post("/v1/gate", json=_action_payload(port=443)) + second = client.post("/v1/gate", json=_action_payload(port=443)).get_json() + third_response = client.post("/v1/gate", json=_action_payload(port=443)) + third = third_response.get_json() + + stored_third = api._decision_history[-1][0] + assert stored_third.id == third["trace_id"] + assert stored_third.similar_decision_ids == third["similar_decision_ids"] + if third["reasons"]: + assert stored_third.risk_flags == third["reasons"] + assert stored_third.risk_flags != [] + + # The third, near-identical action should find at least the second as a match. + assert third["similar_decision_ids"] != [] + assert second["trace_id"] in third["similar_decision_ids"] + + +def test_recorded_decision_carries_the_real_verdict(client) -> None: + """TraceDecision.verdict must reflect the actual gate verdict, not be left + at its dataclass default and reconstructed later from a hardcoded score cutoff.""" + response = client.post("/v1/gate", json=_action_payload(port=443)) + body = response.get_json() + + stored = api._decision_history[-1][0] + assert stored.verdict == body["verdict"] + assert stored.verdict != "" + + +def test_decision_history_per_agent_cap_protects_a_quiet_agent(client) -> None: + """A noisy agent flooding the gate must not evict a quiet agent's decision + history entirely -- see api._DECISION_HISTORY_PER_AGENT_CAP.""" + quiet_payload = _action_payload(agent_id="quiet-agent", target="fw-corp-https", port=443) + client.post("/v1/gate", json=quiet_payload) + + noisy_payload = _action_payload(agent_id="noisy-agent", target="fw-corp-https", port=443) + for _ in range(250): # past both the per-agent sub-cap and the 200-entry total cap + client.post("/v1/gate", json=noisy_payload) + + agent_ids = {decision.agent_id for decision, _vec in api._decision_history} + assert "quiet-agent" in agent_ids + noisy_count = sum( + 1 for decision, _vec in api._decision_history if decision.agent_id == "noisy-agent" + ) + assert noisy_count <= 40 + + +def test_repeated_refused_action_scores_at_least_as_high_the_second_time(client) -> None: + """End-to-end repeat-offense signal through the live /v1/gate handler.""" + payload = _action_payload(agent_id="ghost-agent", target="fw-restricted") + first = client.post("/v1/gate", json=payload).get_json() + second = client.post("/v1/gate", json=payload).get_json() + + assert first["verdict"] in {"WOULD-BLOCK", "BLOCK"} + assert second["score"] >= first["score"] + assert any(first["trace_id"] in reason for reason in second["reasons"]) + + +def test_offense_memory_persists_across_a_simulated_restart(client, tmp_path, monkeypatch) -> None: + """The durability requirement: block an agent, restart the process, confirm + the next similar action still scores higher because of the earlier block.""" + storage = tmp_path / "offense-memory.json" + monkeypatch.setenv("DUSK_OFFENSE_MEMORY_PATH", str(storage)) + reset_config() + api.reset_gate_engine() + + payload = _action_payload(agent_id="ghost-agent", target="fw-restricted") + before_restart = client.post("/v1/gate", json=payload).get_json() + assert before_restart["verdict"] in {"WOULD-BLOCK", "BLOCK"} + # The write lands on a background thread; wait for it before asserting on disk. + api._get_gate_engine().offense_memory.flush() + assert storage.exists() + + # Simulate a process restart: drop the cached engine, force a fresh load + # from disk, exactly like a new process would. + api.reset_gate_engine() + + after_restart = client.post("/v1/gate", json=payload).get_json() + assert after_restart["score"] >= before_restart["score"] + assert any(before_restart["trace_id"] in reason for reason in after_restart["reasons"]) diff --git a/examples/agent-action-monitor/tests/test_actions_azure.py b/examples/agent-action-monitor/tests/test_actions_azure.py new file mode 100644 index 000000000..b18668f48 --- /dev/null +++ b/examples/agent-action-monitor/tests/test_actions_azure.py @@ -0,0 +1,85 @@ +"""Tests for the Azure activity-log adapter.""" + +from __future__ import annotations + +import pytest + +from dusk.actions.adapters.azure import AzureAdapter +from dusk.actions.adapters.base import AdapterError + + +def _nsg_record(**overrides: object) -> dict[str, object]: + """Build a sample Azure NSG security-rule activity-log record.""" + base: dict[str, object] = { + "operationName": { + "value": "Microsoft.Network/networkSecurityGroups/securityRules/write", + "localizedValue": "Create or update security rule", + }, + "caller": "netops-agent@example.com", + "resourceId": ( + "/subscriptions/sub-1/resourceGroups/rg-net/providers/" + "Microsoft.Network/networkSecurityGroups/nsg-corp/securityRules/allow-https" + ), + "eventTimestamp": "2023-11-14T22:13:20Z", + "correlationId": "corr-123", + "properties": {"before": None, "after": {"port": 443}}, + } + base.update(overrides) + return base + + +def test_azure_nsg_maps_to_firewall_rule_change() -> None: + """An NSG security-rule write maps to firewall_rule_change.""" + action = AzureAdapter().parse(_nsg_record()) + assert action.action_type == "firewall_rule_change" + assert action.agent_id == "netops-agent@example.com" + assert action.target.endswith("securityRules/allow-https") + assert action.source == "azure" + assert action.raw_ref == "corr-123" + assert action.change["after"] == {"port": 443} + assert action.timestamp.tzinfo is not None + + +def test_azure_route_table_maps_to_route_change() -> None: + """A routeTables operation maps to route_change.""" + record = _nsg_record( + operationName="Microsoft.Network/routeTables/routes/write", + resourceId="/subscriptions/sub-1/.../routeTables/rt-corp/routes/default", + ) + assert AzureAdapter().parse(record).action_type == "route_change" + + +def test_azure_role_assignment_maps() -> None: + """A roleAssignments operation maps to role_assignment.""" + record = _nsg_record(operationName="Microsoft.Authorization/roleAssignments/write") + assert AzureAdapter().parse(record).action_type == "role_assignment" + + +def test_azure_unknown_operation_maps_to_unknown() -> None: + """An unrecognised operation falls back to the unknown action_type.""" + record = _nsg_record(operationName="Microsoft.Storage/storageAccounts/write") + assert AzureAdapter().parse(record).action_type == "unknown" + + +def test_azure_uses_authorization_action_fallback() -> None: + """When operationName is absent, the RBAC authorization action is used.""" + record = _nsg_record() + del record["operationName"] + record["authorization"] = {"action": "Microsoft.Network/routeTables/write"} + assert AzureAdapter().parse(record).action_type == "route_change" + + +def test_azure_missing_caller_raises_adapter_error() -> None: + """A record with no identity cannot yield a valid AgentAction.""" + record = _nsg_record() + del record["caller"] + with pytest.raises(AdapterError): + AzureAdapter().parse(record) + + +def test_azure_missing_timestamp_raises_adapter_error() -> None: + """A record with no timestamp cannot yield a valid AgentAction.""" + record = _nsg_record() + del record["eventTimestamp"] + with pytest.raises(AdapterError): + AzureAdapter().parse(record) diff --git a/examples/agent-action-monitor/tests/test_actions_bedrock.py b/examples/agent-action-monitor/tests/test_actions_bedrock.py new file mode 100644 index 000000000..849df4bb0 --- /dev/null +++ b/examples/agent-action-monitor/tests/test_actions_bedrock.py @@ -0,0 +1,103 @@ +"""Tests for the Bedrock tool-call adapter.""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest + +from dusk.actions.adapters.base import AdapterError +from dusk.actions.adapters.bedrock import BedrockAdapter + + +def _firewall_tool_use(**overrides: object) -> dict[str, object]: + """Build a sample Bedrock toolUse block proposing a firewall change.""" + base: dict[str, object] = { + "toolUseId": "tooluse-abc123", + "name": "update_firewall_rule", + "input": { + "target": "fw-corp-restricted-segment", + "before": None, + "after": {"port": 22, "cidr": "0.0.0.0/0"}, + }, + } + base.update(overrides) + return base + + +def test_firewall_tool_name_maps_to_firewall_rule_change() -> None: + """A tool name containing 'firewall' maps to firewall_rule_change.""" + action = BedrockAdapter().parse_tool_use( + _firewall_tool_use(), + agent_id="ops-agent-1", + timestamp=datetime(2026, 7, 10, tzinfo=UTC), + ) + assert action.action_type == "firewall_rule_change" + assert action.agent_id == "ops-agent-1" + assert action.target == "fw-corp-restricted-segment" + assert action.source == "bedrock" + assert action.raw_ref == "tooluse-abc123" + assert action.change["after"] == {"port": 22, "cidr": "0.0.0.0/0"} + assert action.timestamp.tzinfo is not None + + +def test_route_tool_name_maps_to_route_change() -> None: + """A tool name containing 'route' maps to route_change.""" + tool_use = _firewall_tool_use(name="update_route_table", input={"target": "rt-1"}) + action = BedrockAdapter().parse_tool_use( + tool_use, agent_id="ops-agent-1", timestamp=datetime(2026, 7, 10, tzinfo=UTC) + ) + assert action.action_type == "route_change" + + +def test_unrecognised_tool_name_maps_to_unknown() -> None: + """A tool name that matches no rule maps to unknown, not an error.""" + tool_use = _firewall_tool_use(name="get_weather", input={"target": "n/a"}) + action = BedrockAdapter().parse_tool_use( + tool_use, agent_id="ops-agent-1", timestamp=datetime(2026, 7, 10, tzinfo=UTC) + ) + assert action.action_type == "unknown" + + +def test_missing_target_raises_adapter_error() -> None: + """A toolUse input with no target cannot yield a valid AgentAction.""" + tool_use = _firewall_tool_use(input={"before": None, "after": {}}) + with pytest.raises(AdapterError, match="target"): + BedrockAdapter().parse_tool_use( + tool_use, agent_id="ops-agent-1", timestamp=datetime(2026, 7, 10, tzinfo=UTC) + ) + + +def test_missing_input_raises_adapter_error() -> None: + """A toolUse block with no input block at all is rejected.""" + tool_use = {"toolUseId": "tooluse-x", "name": "update_firewall_rule"} + with pytest.raises(AdapterError, match="input"): + BedrockAdapter().parse_tool_use( + tool_use, agent_id="ops-agent-1", timestamp=datetime(2026, 7, 10, tzinfo=UTC) + ) + + +def test_parse_satisfies_source_adapter_contract() -> None: + """The registry-facing parse() accepts a raw dict with agent_id/timestamp attached.""" + raw = { + "tool_use": _firewall_tool_use(), + "agent_id": "ops-agent-1", + "timestamp": "2026-07-10T00:00:00+00:00", + } + action = BedrockAdapter().parse(raw) + assert action.action_type == "firewall_rule_change" + assert action.agent_id == "ops-agent-1" + + +def test_parse_missing_tool_use_raises() -> None: + """parse() rejects a raw record with no tool_use block.""" + with pytest.raises(AdapterError, match="tool_use"): + BedrockAdapter().parse({"agent_id": "a", "timestamp": "2026-07-10T00:00:00+00:00"}) + + +def test_registered_in_normaliser() -> None: + """BedrockAdapter is registered under the 'bedrock' source name.""" + from dusk.actions.normaliser import get_adapter, known_sources + + assert "bedrock" in known_sources() + assert isinstance(get_adapter("bedrock"), BedrockAdapter) diff --git a/examples/agent-action-monitor/tests/test_actions_event.py b/examples/agent-action-monitor/tests/test_actions_event.py new file mode 100644 index 000000000..0dcda8af4 --- /dev/null +++ b/examples/agent-action-monitor/tests/test_actions_event.py @@ -0,0 +1,86 @@ +"""Tests for the AgentAction schema: construction, validation, round-trip.""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest + +from dusk.actions.event import AgentAction + +_TS = datetime(2023, 11, 14, 22, 13, 20, tzinfo=UTC) + + +def _action(**overrides: object) -> AgentAction: + """Build a valid AgentAction, overriding fields as needed.""" + kwargs: dict[str, object] = { + "agent_id": "netops-agent", + "timestamp": _TS, + "action_type": "firewall_rule_change", + "target": "fw-corp-https", + "change": {"before": None, "after": {"port": 443}}, + "source": "generic", + "raw_ref": "evt-1", + } + kwargs.update(overrides) + return AgentAction(**kwargs) # type: ignore[arg-type] + + +def test_valid_agent_action_builds() -> None: + """A well-formed AgentAction constructs and exposes its fields.""" + action = _action() + assert action.agent_id == "netops-agent" + assert action.action_type == "firewall_rule_change" + assert action.target == "fw-corp-https" + assert action.change["after"] == {"port": 443} + + +def test_empty_agent_id_raises() -> None: + """An empty or whitespace agent_id is rejected.""" + with pytest.raises(ValueError, match="agent_id"): + _action(agent_id=" ") + + +def test_empty_target_raises() -> None: + """An empty target is rejected.""" + with pytest.raises(ValueError, match="target"): + _action(target="") + + +def test_naive_timestamp_raises() -> None: + """A timezone-naive timestamp is rejected, not silently coerced.""" + with pytest.raises(ValueError, match="timezone-aware"): + _action(timestamp=datetime(2023, 11, 14, 22, 13, 20)) + + +def test_unknown_action_type_raises() -> None: + """An action_type outside the known set is rejected.""" + with pytest.raises(ValueError, match="action_type"): + _action(action_type="totally_made_up") + + +def test_unknown_is_a_valid_action_type() -> None: + """The explicit 'unknown' verb is accepted.""" + assert _action(action_type="unknown").action_type == "unknown" + + +def test_to_dict_is_json_safe() -> None: + """to_dict renders the timestamp as an ISO 8601 string.""" + payload = _action().to_dict() + assert payload["timestamp"] == _TS.isoformat() + assert payload["action_type"] == "firewall_rule_change" + + +def test_to_from_dict_round_trips() -> None: + """from_dict(to_dict(x)) reproduces the original exactly.""" + original = _action() + restored = AgentAction.from_dict(original.to_dict()) + assert restored == original + + +def test_from_dict_missing_field_raises() -> None: + """from_dict rejects a mapping that is missing a required field.""" + payload = _action().to_dict() + del payload["target"] + with pytest.raises(ValueError, match="target"): + AgentAction.from_dict(payload) diff --git a/examples/agent-action-monitor/tests/test_actions_gate.py b/examples/agent-action-monitor/tests/test_actions_gate.py new file mode 100644 index 000000000..eeaf94f91 --- /dev/null +++ b/examples/agent-action-monitor/tests/test_actions_gate.py @@ -0,0 +1,545 @@ +"""Tests for the agent action gate: baseline, analyse, verdict, and benchmark.""" + +from __future__ import annotations + +import os +import sys +from datetime import UTC, datetime + +LAB_DIR = os.path.join(os.path.dirname(__file__), "..", "lab", "actions") +sys.path.insert(0, os.path.abspath(LAB_DIR)) + +import generate_actions # noqa: E402 + +from dusk.actions import baseline as baseline_module # noqa: E402 +from dusk.actions.analyse import analyse # noqa: E402 +from dusk.actions.baseline import Baseline, target_class # noqa: E402 +from dusk.actions.event import AgentAction # noqa: E402 +from dusk.actions.normaliser import normalise_record # noqa: E402 +from dusk.actions.offense_memory import OffenseMemory, OffenseRecord # noqa: E402 +from dusk.actions.verdict import ALLOW, BLOCK, WOULD_BLOCK, ActionGate # noqa: E402 +from dusk.config import Config # noqa: E402 +from dusk.trace.vector import ExtractedTerm # noqa: E402 + +CONFIG = Config() +_TS = datetime(2023, 11, 14, 22, 13, 20, tzinfo=UTC) + + +def _action(agent_id: str, action_type: str, target: str, **change: object) -> AgentAction: + return AgentAction( + agent_id=agent_id, + timestamp=_TS, + action_type=action_type, + target=target, + change={"before": None, "after": dict(change) if change else None}, + source="generic", + ) + + +def _normal() -> list[AgentAction]: + return [normalise_record("generic", r) for r in generate_actions.normal_actions()] + + +def _attacks() -> list[AgentAction]: + return [normalise_record("generic", r) for r in generate_actions.out_of_pattern_actions()] + + +# --- baseline ---------------------------------------------------------------- + + +def test_baseline_learns_per_agent() -> None: + """Learning records one profile per distinct agent.""" + baseline = Baseline.learn(_normal()) + assert set(baseline.agents) == {"netops-agent", "segment-agent", "iam-agent"} + netops = baseline.profile_for("netops-agent") + assert netops is not None + assert "firewall_rule_change" in netops.action_types + assert "fw" in netops.target_classes + + +def test_target_class_groups_by_first_token() -> None: + """Targets that share a prefix share a class.""" + assert target_class("fw-corp-https") == "fw" + assert target_class("fw-guest-to-restricted") == "fw" + assert target_class("seg-corporate") == "seg" + + +def test_change_values_flattens_nested_dicts() -> None: + """A value buried in a nested dict must not be invisible to scoring.""" + change = {"before": None, "after": {"rules": {"cidr": "0.0.0.0/0", "port": 22}}} + values = baseline_module._change_values(change) + assert "0.0.0.0/0" in values + assert "22" in values + + +def test_change_values_flattens_nested_lists() -> None: + """A value buried inside a list of dicts must not be invisible to scoring.""" + change = { + "before": None, + "after": {"rules": [{"cidr": "10.0.0.0/8"}, {"cidr": "0.0.0.0/0", "port": 22}]}, + } + values = baseline_module._change_values(change) + assert "0.0.0.0/0" in values + assert "10.0.0.0/8" in values + assert "22" in values + + +def test_change_values_still_flattens_top_level() -> None: + """The original flat-dict behaviour is unchanged.""" + change = {"before": None, "after": {"port": 443}} + assert baseline_module._change_values(change) == {"443"} + + +def test_change_values_depth_is_bounded() -> None: + """An adversarially deep payload does not make flattening unbounded.""" + nested: dict[str, object] = {"leaf": "0.0.0.0/0"} + for _ in range(20): + nested = {"wrapper": nested} + change = {"before": None, "after": nested} + # Should not raise (e.g. RecursionError) and should not necessarily find + # the deeply buried leaf, since depth is capped defensively. + baseline_module._change_values(change) + + +# --- analyse ----------------------------------------------------------------- + + +def test_known_good_action_scores_zero() -> None: + """An action matching the agent's baseline is not anomalous.""" + baseline = Baseline.learn(_normal()) + result = analyse(baseline, _action("netops-agent", "firewall_rule_change", "fw-corp-https")) + assert result.score == 0.0 + assert result.reasons == ["action matches the agent's established pattern"] + + +def test_new_action_type_is_flagged() -> None: + """An agent doing a verb it never does scores high and maps to ATT&CK.""" + baseline = Baseline.learn(_normal()) + result = analyse(baseline, _action("segment-agent", "firewall_rule_change", "fw-x")) + assert result.score >= CONFIG.gate_block_threshold + assert "T1562" in result.mitre_attack + assert result.mitre_atlas.startswith("AML.") + + +def test_privilege_escalation_is_flagged() -> None: + """Granting a sensitive role is caught even when the verb is familiar.""" + baseline = Baseline.learn(_normal()) + result = analyse(baseline, _action("iam-agent", "role_assignment", "ra-self", role="owner")) + assert result.score >= CONFIG.gate_block_threshold + assert result.blast_radius == "high" + assert any("sensitive" in r for r in result.reasons) + + +def test_nested_privilege_escalation_is_flagged() -> None: + """A sensitive value buried in a nested change payload is not invisible. + + Same scenario as test_privilege_escalation_is_flagged, but the sensitive + value sits inside a nested structure -- realistic for a control-plane + payload shaped like {"after": {"rules": [{"role": "owner"}]}}. Before the + nested-flatten fix, this would score 0.0 and pass through silently. + """ + baseline = Baseline.learn(_normal()) + action = AgentAction( + agent_id="iam-agent", + timestamp=_TS, + action_type="role_assignment", + target="ra-self", + change={"before": None, "after": {"grants": [{"role": "owner", "scope": "global"}]}}, + source="generic", + ) + result = analyse(baseline, action) + assert result.score >= CONFIG.gate_block_threshold + assert result.blast_radius == "high" + assert any("sensitive" in r for r in result.reasons) + + +def test_unknown_agent_is_noted() -> None: + """An agent with no baseline is called out.""" + baseline = Baseline.learn(_normal()) + result = analyse(baseline, _action("ghost-agent", "route_change", "rt-x")) + assert result.score > 0.0 + assert any("no established baseline" in r for r in result.reasons) + + +def test_agent_history_without_sie_does_not_change_score() -> None: + """Passing history with SIE unavailable is a no-op (the default, no-SIE case).""" + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-corp-https", port=443) + with_history = analyse(baseline, action, agent_history=_normal()) + without_history = analyse(baseline, action) + assert with_history.score == without_history.score + assert with_history.reasons == without_history.reasons + + +def test_agent_history_low_rerank_similarity_adds_reason() -> None: + """A low SIE rerank score adds a reason and raises the score, on top of rule-based checks.""" + from unittest.mock import patch + + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-corp-https", port=443) + baseline_only = analyse(baseline, action) + + with patch("dusk.actions.analyse.sie_score", return_value=[0.05, 0.05]): + reranked = analyse(baseline, action, agent_history=_normal()[:2]) + + assert reranked.score > baseline_only.score + assert any("SIE rerank" in r for r in reranked.reasons) + + +def test_agent_history_high_rerank_similarity_is_unchanged() -> None: + """A confident rerank match does not add the low-similarity reason.""" + from unittest.mock import patch + + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-corp-https", port=443) + baseline_only = analyse(baseline, action) + + with patch("dusk.actions.analyse.sie_score", return_value=[0.9, 0.9]): + reranked = analyse(baseline, action, agent_history=_normal()[:2]) + + assert reranked.score == baseline_only.score + assert not any("SIE rerank" in r for r in reranked.reasons) + + +def test_sie_extract_flags_terms_missed_by_the_static_frozenset() -> None: + """A GLiNER-extracted term outside the hardcoded sensitive set adds a reason.""" + from unittest.mock import patch + + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-corp-https", port=443) + baseline_only = analyse(baseline, action) + + term = ExtractedTerm(text="superuser", label="role", score=0.95) + with patch("dusk.actions.analyse.sie_extract", return_value=[term]): + extracted = analyse(baseline, action) + + assert extracted.score > baseline_only.score + assert any("SIE extract" in r and "superuser" in r for r in extracted.reasons) + + +def test_sie_extract_terms_already_in_static_set_are_not_duplicated() -> None: + """A term the static frozenset already catches doesn't add a second reason.""" + from unittest.mock import patch + + baseline = Baseline.learn(_normal()) + action = _action("iam-agent", "role_assignment", "ra-self", role="owner") + baseline_only = analyse(baseline, action) + + term = ExtractedTerm(text="owner", label="role", score=0.95) + with patch("dusk.actions.analyse.sie_extract", return_value=[term]): + extracted = analyse(baseline, action) + + assert extracted.score == baseline_only.score + assert not any("SIE extract" in r for r in extracted.reasons) + + +def test_sie_extract_low_confidence_term_is_ignored() -> None: + """A GLiNER hit below the confidence floor is dropped, not counted as evidence.""" + from unittest.mock import patch + + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-corp-https", port=443) + baseline_only = analyse(baseline, action) + + term = ExtractedTerm(text="superuser", label="role", score=0.2) + with patch("dusk.actions.analyse.sie_extract", return_value=[term]): + extracted = analyse(baseline, action) + + assert extracted.score == baseline_only.score + + +def test_sie_extract_high_confidence_non_privileged_label_is_ignored() -> None: + """A confident but ordinary resource/port extraction is not privilege escalation.""" + from unittest.mock import patch + + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-corp-https", port=443) + baseline_only = analyse(baseline, action) + + terms = [ + ExtractedTerm(text="eth0", label="resource", score=0.95), + ExtractedTerm(text="8080", label="port", score=0.95), + ] + with patch("dusk.actions.analyse.sie_extract", return_value=terms): + extracted = analyse(baseline, action) + + assert extracted.score == baseline_only.score + assert not any("SIE extract" in r for r in extracted.reasons) + assert not any("SIE extract" in r for r in extracted.reasons) + + +def test_sie_extract_unavailable_does_not_change_score() -> None: + """The default (no-SIE) case: sie_extract returns [] and nothing changes.""" + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-corp-https", port=443) + result = analyse(baseline, action) + assert not any("SIE extract" in r for r in result.reasons) + + +# --- repeat-offense signal ----------------------------------------------------- + + +def _offense(**overrides: object) -> OffenseRecord: + defaults: dict[str, object] = { + "trace_id": "trace-offense-1", + "agent_id": "netops-agent", + "action_type": "firewall_rule_change", + "target_class": "fw", + "tokens": ("fw", "restricted"), + "verdict": "BLOCK", + "timestamp": datetime.now(UTC), + } + defaults.update(overrides) + return OffenseRecord(**defaults) # type: ignore[arg-type] + + +def test_no_offenses_does_not_change_score() -> None: + """A clean-history agent is completely unaffected by the repeat-offense signal.""" + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-restricted") + without = analyse(baseline, action) + with_empty = analyse(baseline, action, offenses=[]) + assert without.score == with_empty.score + assert without.reasons == with_empty.reasons + + +def test_matching_offense_raises_score_and_cites_trace_id() -> None: + """A same-type, same-target-class repeat past offense adds score and names the prior trace.""" + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-restricted") + baseline_only = analyse(baseline, action) + offense = _offense(trace_id="trace-xyz") + + with_offense = analyse(baseline, action, offenses=[offense], config=CONFIG) + + assert with_offense.score > baseline_only.score + assert any("trace-xyz" in r for r in with_offense.reasons) + + +def test_different_action_type_offense_does_not_match() -> None: + """An offense for a different action type must not contribute -- type match is required.""" + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "role_assignment", "ra-self", role="owner") + baseline_only = analyse(baseline, action) + offense = _offense(action_type="firewall_rule_change") + + with_offense = analyse(baseline, action, offenses=[offense], config=CONFIG) + + assert with_offense.score == baseline_only.score + + +def test_repeat_offense_contribution_is_capped() -> None: + """The signal alone cannot exceed repeat_offense_max_contribution, however strong the match.""" + config = Config(repeat_offense_max_contribution=0.05) + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-restricted") + baseline_only = analyse(baseline, action) + offense = _offense() + + with_offense = analyse(baseline, action, offenses=[offense], config=config) + + assert with_offense.score - baseline_only.score <= 0.05 + 1e-9 + + +def test_old_offense_contributes_less_than_a_recent_one() -> None: + """Decay: an offense from long ago must weigh less than one from moments ago.""" + from datetime import timedelta + + config = Config(repeat_offense_half_life_days=10.0) + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-restricted") + + recent = analyse( + baseline, action, offenses=[_offense(timestamp=datetime.now(UTC))], config=config + ) + old = analyse( + baseline, + action, + offenses=[_offense(timestamp=datetime.now(UTC) - timedelta(days=100))], + config=config, + ) + + assert recent.score > old.score + + +def test_multiple_offenses_use_the_single_best_match_not_the_sum() -> None: + """Anti-gaming: flooding with many weak matches must not out-score one strong match.""" + config = Config(repeat_offense_max_contribution=1.0) + baseline = Baseline.learn(_normal()) + # A known action/target for this agent, so the deterministic checks below + # contribute 0 and only the repeat-offense signal moves the score -- + # otherwise both cases would saturate at the 1.0 clamp and be indistinguishable. + action = _action("netops-agent", "firewall_rule_change", "fw-corp-https", port=443) + + single_strong = analyse( + baseline, + action, + offenses=[_offense(target_class="fw", tokens=("fw", "corp", "https"))], + config=config, + ) + many_weak = analyse( + baseline, + action, + offenses=[_offense(target_class="seg", tokens=("seg",)) for _ in range(20)], + config=config, + ) + + # The weak matches don't even share a target class or token, so they + # contribute nothing at all -- confirming there is no additive stacking. + assert many_weak.score < single_strong.score + + +def test_offense_reason_names_the_verdict_and_date() -> None: + baseline = Baseline.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-restricted") + offense = _offense(verdict="WOULD-BLOCK") + + result = analyse(baseline, action, offenses=[offense], config=CONFIG) + + assert any("WOULD-BLOCK" in r for r in result.reasons) + + +# --- verdict ----------------------------------------------------------------- + + +def test_gate_allows_routine_refuses_attacks() -> None: + """Watch mode allows routine actions and WOULD-BLOCKs the attacks.""" + gate = ActionGate(config=CONFIG) + gate.learn(_normal()) + for action in _normal(): + assert gate.evaluate(action).verdict == ALLOW + for attack in _attacks(): + v = gate.evaluate(attack) + assert v.verdict == WOULD_BLOCK + assert v.refused is True + + +def test_gate_learn_tracks_raw_history_per_agent() -> None: + """learn() keeps the raw actions ActionGate.evaluate() feeds into analyse().""" + gate = ActionGate(config=CONFIG) + gate.learn(_normal()) + netops_history = gate._history.get("netops-agent", []) + assert netops_history + assert all(a.agent_id == "netops-agent" for a in netops_history) + + +def test_gate_evaluate_passes_history_to_sie_rerank() -> None: + """A gate evaluation surfaces the SIE rerank reason when the mocked score is low.""" + from unittest.mock import patch + + gate = ActionGate(config=CONFIG) + gate.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-corp-https", port=443) + + with patch("dusk.actions.analyse.sie_score", return_value=[0.05] * 20): + verdict = gate.evaluate(action) + + assert any("SIE rerank" in r for r in verdict.analysis.reasons) + + +def test_enforce_mode_blocks() -> None: + """Enforce mode renders BLOCK instead of WOULD-BLOCK.""" + gate = ActionGate(config=CONFIG, enforce=True) + gate.learn(_normal()) + assert gate.evaluate(_attacks()[0]).verdict == BLOCK + + +def test_empty_baseline_evaluation_does_not_crash() -> None: + """A gate with no baseline still renders a verdict for a single action.""" + gate = ActionGate(config=CONFIG) + verdict = gate.evaluate(_action("a", "route_change", "rt-1")) + assert verdict.verdict in (ALLOW, WOULD_BLOCK) + + +# --- offense memory wiring ----------------------------------------------------- + + +def test_gate_without_offense_memory_behaves_as_before() -> None: + """No offense_memory passed -- the repeat-offense signal is a complete no-op.""" + gate = ActionGate(config=CONFIG) + gate.learn(_normal()) + attack = _attacks()[0] + first = gate.evaluate(attack) + second = gate.evaluate(attack) + assert first.analysis.score == second.analysis.score + + +def test_refused_verdict_is_recorded_in_offense_memory() -> None: + memory = OffenseMemory(storage_path=None) + gate = ActionGate(config=CONFIG, offense_memory=memory) + gate.learn(_normal()) + attack = _attacks()[0] + + verdict = gate.evaluate(attack) + + offenses = memory.offenses_for(attack.agent_id) + assert len(offenses) == 1 + assert offenses[0].trace_id == verdict.trace_id + assert offenses[0].verdict == verdict.verdict + + +def test_allowed_verdict_is_not_recorded_in_offense_memory() -> None: + memory = OffenseMemory(storage_path=None) + gate = ActionGate(config=CONFIG, offense_memory=memory) + gate.learn(_normal()) + + for action in _normal(): + verdict = gate.evaluate(action) + assert verdict.verdict == ALLOW + + assert all(memory.offenses_for(a.agent_id) == [] for a in _normal()) + + +def test_repeated_attack_scores_higher_the_second_time() -> None: + """The end-to-end point of this feature: a repeat offender is judged more harshly.""" + memory = OffenseMemory(storage_path=None) + gate = ActionGate(config=CONFIG, offense_memory=memory) + gate.learn(_normal()) + attack = _attacks()[0] + + first = gate.evaluate(attack) + second = gate.evaluate(attack) + + assert second.analysis.score >= first.analysis.score + assert any(first.trace_id in r for r in second.analysis.reasons) + + +def test_gate_verdict_trace_id_is_unique_per_evaluation() -> None: + gate = ActionGate(config=CONFIG) + gate.learn(_normal()) + action = _action("netops-agent", "firewall_rule_change", "fw-corp-https", port=443) + first = gate.evaluate(action) + second = gate.evaluate(action) + assert first.trace_id != second.trace_id + + +# --- labelled benchmark ------------------------------------------------------ + + +def test_benchmark_precision_recall() -> None: + """On the labelled fixture the gate catches every attack with no false alarms.""" + gate = ActionGate(config=CONFIG) + gate.learn(_normal()) + + labelled = [(a, False) for a in _normal()] + [(a, True) for a in _attacks()] + tp = fp = fn = tn = 0 + for action, is_attack in labelled: + refused = gate.evaluate(action).refused + if is_attack and refused: + tp += 1 + elif is_attack and not refused: + fn += 1 + elif not is_attack and refused: + fp += 1 + else: + tn += 1 + + precision = tp / (tp + fp) if (tp + fp) else 1.0 + recall = tp / (tp + fn) if (tp + fn) else 1.0 + fp_rate = fp / (fp + tn) if (fp + tn) else 0.0 + + # Surface the numbers in the assertion message for the demo record. + assert (precision, recall, fp_rate) == (1.0, 1.0, 0.0), ( + f"precision={precision:.2f} recall={recall:.2f} fp_rate={fp_rate:.2f} " + f"(tp={tp} fp={fp} fn={fn} tn={tn})" + ) diff --git a/examples/agent-action-monitor/tests/test_actions_generic.py b/examples/agent-action-monitor/tests/test_actions_generic.py new file mode 100644 index 000000000..9c5fd0d79 --- /dev/null +++ b/examples/agent-action-monitor/tests/test_actions_generic.py @@ -0,0 +1,62 @@ +"""Tests for the generic source adapter.""" + +from __future__ import annotations + +import pytest + +from dusk.actions.adapters.base import AdapterError +from dusk.actions.adapters.generic import GenericAdapter + + +def _record(**overrides: object) -> dict[str, object]: + """Build one well-formed generic record, overriding fields as needed.""" + base: dict[str, object] = { + "agent_id": "netops-agent", + "timestamp": "2023-11-14T22:13:20+00:00", + "action_type": "firewall_rule_change", + "target": "fw-corp-https", + "change": {"before": None, "after": {"port": 443}}, + "source": "generic", + "raw_ref": "evt-1", + } + base.update(overrides) + return base + + +def test_generic_builds_agent_action() -> None: + """A well-formed generic record maps to the canonical fields.""" + action = GenericAdapter().parse(_record()) + assert action.agent_id == "netops-agent" + assert action.action_type == "firewall_rule_change" + assert action.target == "fw-corp-https" + assert action.source == "generic" + assert action.change["after"] == {"port": 443} + assert action.timestamp.tzinfo is not None + + +def test_generic_defaults_change_when_absent() -> None: + """A record without a change gets the empty before/after delta.""" + record = _record() + del record["change"] + action = GenericAdapter().parse(record) + assert action.change == {"before": None, "after": None} + + +def test_generic_missing_field_raises_adapter_error() -> None: + """A missing required field raises AdapterError, not a raw error.""" + record = _record() + del record["target"] + with pytest.raises(AdapterError, match="target"): + GenericAdapter().parse(record) + + +def test_generic_empty_identity_raises() -> None: + """An empty agent_id is rejected as an AdapterError.""" + with pytest.raises(AdapterError, match="agent_id"): + GenericAdapter().parse(_record(agent_id="")) + + +def test_generic_naive_timestamp_raises() -> None: + """A timezone-naive timestamp is rejected as an AdapterError.""" + with pytest.raises(AdapterError, match="timezone-aware"): + GenericAdapter().parse(_record(timestamp="2023-11-14T22:13:20")) diff --git a/examples/agent-action-monitor/tests/test_actions_ingest.py b/examples/agent-action-monitor/tests/test_actions_ingest.py new file mode 100644 index 000000000..7a90f4d35 --- /dev/null +++ b/examples/agent-action-monitor/tests/test_actions_ingest.py @@ -0,0 +1,100 @@ +"""Tests for ingest_file: reading a JSON list of records via an adapter.""" + +from __future__ import annotations + +import json +import os +import sys +from pathlib import Path + +import pytest + +LAB_DIR = os.path.join(os.path.dirname(__file__), "..", "lab", "actions") +sys.path.insert(0, os.path.abspath(LAB_DIR)) + +import generate_actions # noqa: E402 + +from dusk.actions.event import AgentAction # noqa: E402 +from dusk.actions.ingest import ingest_file # noqa: E402 + +FIXTURES = os.path.join(os.path.dirname(__file__), "fixtures") + + +def _ensure_fixtures() -> tuple[str, str]: + """Return the (normal, mixed) fixture paths, generating them if absent.""" + normal = os.path.join(FIXTURES, "actions_normal.json") + mixed = os.path.join(FIXTURES, "actions_mixed.json") + if not (os.path.exists(normal) and os.path.exists(mixed)): + generate_actions.generate(FIXTURES) + return normal, mixed + + +def _record(**overrides: object) -> dict[str, object]: + """Build one valid generic record.""" + base: dict[str, object] = { + "agent_id": "a", + "timestamp": "2023-11-14T22:13:20+00:00", + "action_type": "route_change", + "target": "rt-1", + "change": {"before": None, "after": None}, + "source": "generic", + } + base.update(overrides) + return base + + +def test_ingest_normal_fixture() -> None: + """The normal fixture ingests every record as an AgentAction.""" + normal, _ = _ensure_fixtures() + actions = ingest_file(normal, "generic") + assert len(actions) == 15 + assert all(isinstance(a, AgentAction) for a in actions) + + +def test_ingest_mixed_fixture() -> None: + """The mixed fixture ingests the routine plus out-of-pattern actions.""" + _, mixed = _ensure_fixtures() + assert len(ingest_file(mixed, "generic")) == 18 + + +def test_ingest_skips_one_malformed_record(tmp_path: Path) -> None: + """One malformed record is skipped; the valid ones still return.""" + records = [_record(), {"agent_id": "a"}] # second is missing fields + path = tmp_path / "actions.json" + path.write_text(json.dumps(records), encoding="utf-8") + assert len(ingest_file(str(path), "generic")) == 1 + + +def test_ingest_empty_list(tmp_path: Path) -> None: + """An empty JSON list yields no actions and does not raise.""" + path = tmp_path / "empty.json" + path.write_text("[]", encoding="utf-8") + assert ingest_file(str(path), "generic") == [] + + +def test_ingest_single_record(tmp_path: Path) -> None: + """A single valid record yields one action.""" + path = tmp_path / "one.json" + path.write_text(json.dumps([_record()]), encoding="utf-8") + assert len(ingest_file(str(path), "generic")) == 1 + + +def test_ingest_missing_file_raises() -> None: + """A missing file raises FileNotFoundError.""" + with pytest.raises(FileNotFoundError): + ingest_file("does-not-exist.json", "generic") + + +def test_ingest_not_a_list_raises(tmp_path: Path) -> None: + """A top-level JSON object (not a list) raises ValueError.""" + path = tmp_path / "object.json" + path.write_text('{"agent_id": "a"}', encoding="utf-8") + with pytest.raises(ValueError, match="must contain a JSON list"): + ingest_file(str(path), "generic") + + +def test_ingest_unknown_source_skips_all(tmp_path: Path) -> None: + """An unknown source means every record is skipped, returning empty.""" + path = tmp_path / "actions.json" + path.write_text(json.dumps([_record()]), encoding="utf-8") + assert ingest_file(str(path), "nope") == [] diff --git a/examples/agent-action-monitor/tests/test_actions_offense_memory.py b/examples/agent-action-monitor/tests/test_actions_offense_memory.py new file mode 100644 index 000000000..ac7f538ee --- /dev/null +++ b/examples/agent-action-monitor/tests/test_actions_offense_memory.py @@ -0,0 +1,281 @@ +"""Tests for OffenseMemory: persistence, per-agent scoping, capping.""" + +from __future__ import annotations + +from pathlib import Path + +from dusk.actions.offense_memory import OffenseMemory, OffenseRecord + + +def _record(memory: OffenseMemory, agent_id: str = "netops-agent", **overrides: object) -> None: + defaults: dict[str, object] = { + "trace_id": "trace-1", + "agent_id": agent_id, + "action_type": "firewall_rule_change", + "target_class": "fw", + "tokens": {"fw", "restricted"}, + "verdict": "BLOCK", + } + defaults.update(overrides) + memory.record(**defaults) # type: ignore[arg-type] + + +def test_in_memory_only_when_no_storage_path() -> None: + memory = OffenseMemory(storage_path=None) + _record(memory) + assert len(memory.offenses_for("netops-agent")) == 1 + + +def test_record_then_offenses_for_returns_it() -> None: + memory = OffenseMemory(storage_path=None) + _record(memory, trace_id="abc123") + offenses = memory.offenses_for("netops-agent") + assert len(offenses) == 1 + assert offenses[0].trace_id == "abc123" + assert offenses[0].verdict == "BLOCK" + + +def test_offenses_are_scoped_per_agent() -> None: + """One agent's offenses must never leak into another agent's list.""" + memory = OffenseMemory(storage_path=None) + _record(memory, agent_id="agent-a") + _record(memory, agent_id="agent-b") + assert len(memory.offenses_for("agent-a")) == 1 + assert len(memory.offenses_for("agent-b")) == 1 + assert memory.offenses_for("agent-c") == [] + + +def test_noisy_agent_does_not_crowd_out_a_quiet_agent() -> None: + """A per-agent cap must not be shared -- a busy agent must not evict a quiet agent's history.""" + memory = OffenseMemory(storage_path=None) + _record(memory, agent_id="quiet-agent", trace_id="quiet-1") + for i in range(100): + _record(memory, agent_id="noisy-agent", trace_id=f"noisy-{i}") + assert len(memory.offenses_for("quiet-agent")) == 1 + assert memory.offenses_for("quiet-agent")[0].trace_id == "quiet-1" + + +def test_per_agent_cap_evicts_oldest_first() -> None: + memory = OffenseMemory(storage_path=None) + for i in range(60): + _record(memory, trace_id=f"trace-{i}") + offenses = memory.offenses_for("netops-agent") + assert len(offenses) == 50 + assert offenses[0].trace_id == "trace-10" + assert offenses[-1].trace_id == "trace-59" + + +def test_persists_across_new_instances(tmp_path: Path) -> None: + """The core durability requirement: memory must survive a process restart.""" + storage = tmp_path / "offenses.json" + first = OffenseMemory(storage_path=str(storage)) + _record(first, trace_id="persisted-1") + first.flush() + assert first.last_persist_error is None + + second = OffenseMemory(storage_path=str(storage)) + offenses = second.offenses_for("netops-agent") + assert len(offenses) == 1 + assert offenses[0].trace_id == "persisted-1" + + +def test_tracked_agent_cap_evicts_the_least_recently_touched_agent() -> None: + """Bounded across agents too, not just within one -- see the class docstring.""" + memory = OffenseMemory(storage_path=None) + for i in range(500): + _record(memory, agent_id=f"agent-{i}", trace_id="t") + + # Touch agent-0 again so it becomes most-recently-used and should survive + # the next eviction, even though it was inserted first. + _record(memory, agent_id="agent-0", trace_id="t2") + # One more distinct agent pushes the tracked-agent count over the cap. + _record(memory, agent_id="agent-500", trace_id="t") + + assert memory.offenses_for("agent-0") != [], "recently-touched agent should survive" + assert memory.offenses_for("agent-1") == [], "true least-recently-touched agent, evicted" + assert memory.offenses_for("agent-500") != [] + + +def test_missing_storage_file_starts_empty(tmp_path: Path) -> None: + memory = OffenseMemory(storage_path=str(tmp_path / "does-not-exist.json")) + assert memory.offenses_for("netops-agent") == [] + + +def test_corrupt_storage_file_starts_empty_without_raising(tmp_path: Path) -> None: + storage = tmp_path / "offenses.json" + storage.write_text("not valid json{{{", encoding="utf-8") + memory = OffenseMemory(storage_path=str(storage)) + assert memory.offenses_for("netops-agent") == [] + + +def test_storage_file_with_wrong_shape_starts_empty(tmp_path: Path) -> None: + storage = tmp_path / "offenses.json" + storage.write_text('["just", "a", "list"]', encoding="utf-8") + memory = OffenseMemory(storage_path=str(storage)) + assert memory.offenses_for("netops-agent") == [] + + +def test_malformed_individual_record_is_skipped_not_fatal(tmp_path: Path) -> None: + import json + + storage = tmp_path / "offenses.json" + storage.write_text( + json.dumps({"netops-agent": [{"trace_id": "only-a-trace-id"}]}), encoding="utf-8" + ) + memory = OffenseMemory(storage_path=str(storage)) + assert memory.offenses_for("netops-agent") == [] + + +def test_clear_wipes_all_agents(tmp_path: Path) -> None: + storage = tmp_path / "offenses.json" + memory = OffenseMemory(storage_path=str(storage)) + _record(memory, agent_id="agent-a") + _record(memory, agent_id="agent-b") + + memory.clear() + memory.flush() + + assert memory.offenses_for("agent-a") == [] + assert memory.offenses_for("agent-b") == [] + reloaded = OffenseMemory(storage_path=str(storage)) + assert reloaded.offenses_for("agent-a") == [] + + +def test_record_does_not_block_on_a_slow_disk_write(tmp_path: Path, monkeypatch) -> None: + """The point of the background writer: record() returns before the write lands, + even when the write itself is slow.""" + import time + + storage = tmp_path / "offenses.json" + original_replace = Path.replace + + def slow_replace(self: Path, target: object) -> object: + time.sleep(0.2) + return original_replace(self, target) + + monkeypatch.setattr(Path, "replace", slow_replace) + + memory = OffenseMemory(storage_path=str(storage)) + start = time.monotonic() + _record(memory, trace_id="fast-return") + elapsed = time.monotonic() - start + + assert elapsed < 0.2, "record() waited on the disk write instead of backgrounding it" + assert not storage.exists(), "write should still be in flight at this point" + + memory.flush() + assert storage.exists() + + +def test_flush_is_a_no_op_when_nothing_was_ever_written() -> None: + OffenseMemory(storage_path=None).flush() + + +def test_concurrent_first_records_create_only_one_executor(tmp_path: Path, monkeypatch) -> None: + """Many threads racing to be the first to schedule a write must not each + create their own executor -- the lazy-init must be lock-protected.""" + import threading + from concurrent.futures import ThreadPoolExecutor + + created: list[int] = [] + original_init = ThreadPoolExecutor.__init__ + + def counting_init(self: ThreadPoolExecutor, *args: object, **kwargs: object) -> None: + created.append(1) + original_init(self, *args, **kwargs) # type: ignore[arg-type] + + monkeypatch.setattr(ThreadPoolExecutor, "__init__", counting_init) + + storage = tmp_path / "offenses.json" + memory = OffenseMemory(storage_path=str(storage)) + + threads = [ + threading.Thread(target=_record, kwargs={"memory": memory, "trace_id": f"race-{i}"}) + for i in range(50) + ] + for t in threads: + t.start() + for t in threads: + t.join() + memory.flush() + + assert len(created) == 1, f"expected exactly one executor, created {len(created)}" + assert len(memory.offenses_for("netops-agent")) == 50 + + +def test_burst_of_records_during_a_slow_write_coalesces(tmp_path: Path, monkeypatch) -> None: + """A burst of records that arrives while a write is in flight costs at most + a couple of writes, not one write per record, and loses nothing.""" + import time + + storage = tmp_path / "offenses.json" + memory = OffenseMemory(storage_path=str(storage)) + + write_count = 0 + original_write = OffenseMemory._write_to_disk + + def counting_write(self: OffenseMemory, payload: dict[str, object]) -> None: + nonlocal write_count + write_count += 1 + time.sleep(0.05) + original_write(self, payload) + + monkeypatch.setattr(OffenseMemory, "_write_to_disk", counting_write) + + for i in range(20): + _record(memory, trace_id=f"trace-{i}") + + memory.flush() + + assert write_count <= 3, f"expected coalescing, got {write_count} writes for 20 records" + assert len(memory.offenses_for("netops-agent")) == 20 + reloaded = OffenseMemory(storage_path=str(storage)) + assert len(reloaded.offenses_for("netops-agent")) == 20 + + +def test_close_flushes_then_stops_scheduling_new_writes(tmp_path: Path) -> None: + storage = tmp_path / "offenses.json" + memory = OffenseMemory(storage_path=str(storage)) + _record(memory, trace_id="before-close") + memory.close() + assert storage.exists() + + reloaded = OffenseMemory(storage_path=str(storage)) + assert len(reloaded.offenses_for("netops-agent")) == 1 + + # Still updates in-memory state after close, but no longer persists it. + _record(memory, trace_id="after-close") + assert len(memory.offenses_for("netops-agent")) == 2 + reloaded_again = OffenseMemory(storage_path=str(storage)) + assert len(reloaded_again.offenses_for("netops-agent")) == 1 + + memory.close() # safe to call twice + + +def test_last_persist_error_is_set_on_a_failed_write(tmp_path: Path) -> None: + blocker = tmp_path / "not-a-directory" + blocker.write_text("blocking file", encoding="utf-8") + bad_storage = blocker / "offenses.json" + + memory = OffenseMemory(storage_path=str(bad_storage)) + assert memory.last_persist_error is None + _record(memory, trace_id="will-fail") + memory.flush() + + assert memory.last_persist_error is not None + + +def test_offense_record_round_trips_through_to_dict_from_dict() -> None: + from datetime import UTC, datetime + + record = OffenseRecord( + trace_id="t1", + agent_id="netops-agent", + action_type="firewall_rule_change", + target_class="fw", + tokens=("fw", "restricted"), + verdict="BLOCK", + timestamp=datetime(2024, 1, 1, tzinfo=UTC), + ) + rebuilt = OffenseRecord.from_dict(record.to_dict()) + assert rebuilt == record diff --git a/examples/agent-action-monitor/tests/test_config.py b/examples/agent-action-monitor/tests/test_config.py new file mode 100644 index 000000000..a87cf1965 --- /dev/null +++ b/examples/agent-action-monitor/tests/test_config.py @@ -0,0 +1,139 @@ +"""Tests for the configuration system: defaults, YAML, env vars, validation.""" + +from __future__ import annotations + +import os +from collections.abc import Iterator +from pathlib import Path + +import pytest + +from dusk.config import Config, ConfigError, load_config + + +def test_defaults() -> None: + """A default Config exposes the documented threshold values.""" + config = Config() + assert config.sweep_threshold == 15 + assert config.sweep_window_seconds == 10.0 + assert config.sweep_timing_std_threshold == 0.05 + assert config.boundary_port_threshold == 10 + assert config.boundary_window_seconds == 30.0 + assert config.alert_log_path == "dusk-alerts.json" + assert config.log_level == "WARNING" + assert config.enforce is False + assert config.sie_endpoint == "http://sie:8080" + assert config.sie_encode_model == "BAAI/bge-m3" + assert config.sie_score_model == "BAAI/bge-reranker-v2-m3" + assert config.sie_extract_model == "urchade/gliner_multi-v2.1" + assert config.sie_timeout_ms == 10000 + assert config.n8n_alert_url == "" + assert config.n8n_report_url == "" + assert config.n8n_decision_url == "" + assert config.n8n_max_workers == 8 + assert config.n8n_max_queued == 200 + assert config.offense_memory_path == "" + assert config.repeat_offense_max_contribution == 0.3 + assert config.repeat_offense_half_life_days == 30.0 + + +def test_load_defaults_when_no_file(tmp_path: Path) -> None: + """load_config falls back to defaults when no YAML file is present.""" + config = load_config(str(tmp_path / "missing.yaml")) + assert config.sweep_threshold == 15 + + +def test_yaml_overrides(tmp_path: Path) -> None: + """Values in dusk.yaml override the dataclass defaults.""" + yaml_path = tmp_path / "dusk.yaml" + yaml_path.write_text("sweep_threshold: 5\nboundary_port_threshold: 3\n") + config = load_config(str(yaml_path)) + assert config.sweep_threshold == 5 + assert config.boundary_port_threshold == 3 + # Untouched fields keep their defaults. + assert config.sweep_window_seconds == 10.0 + + +@pytest.fixture +def _clean_env() -> Iterator[None]: + """Remove DUSK_* env vars around a test.""" + saved = {k: v for k, v in os.environ.items() if k.startswith("DUSK_")} + for key in saved: + del os.environ[key] + yield + for key in list(os.environ): + if key.startswith("DUSK_"): + del os.environ[key] + os.environ.update(saved) + + +def test_env_overrides_yaml(tmp_path: Path, _clean_env: None) -> None: + """DUSK_* environment variables take precedence over YAML and defaults.""" + yaml_path = tmp_path / "dusk.yaml" + yaml_path.write_text("sweep_threshold: 5\n") + os.environ["DUSK_SWEEP_THRESHOLD"] = "42" + config = load_config(str(yaml_path)) + assert config.sweep_threshold == 42 + + +def test_invalid_threshold_raises() -> None: + """A non-positive threshold is rejected at construction.""" + with pytest.raises(ConfigError): + Config(sweep_threshold=0) + + +def test_invalid_log_level_raises() -> None: + """An unrecognised log level name is rejected.""" + with pytest.raises(ConfigError): + Config(log_level="NONSENSE") + + +def test_invalid_sie_timeout_raises() -> None: + """A non-positive SIE timeout is rejected at construction.""" + with pytest.raises(ConfigError): + Config(sie_timeout_ms=0) + + +@pytest.mark.parametrize( + "field", ["sie_endpoint", "n8n_alert_url", "n8n_report_url", "n8n_decision_url"] +) +def test_invalid_url_scheme_raises(field: str) -> None: + """A configured URL without an http(s) scheme is rejected.""" + with pytest.raises(ConfigError): + Config(**{field: "ftp://example.com"}) + + +def test_empty_n8n_urls_are_valid() -> None: + """An unset (empty) n8n URL is valid -- it just means that webhook is off.""" + config = Config(n8n_alert_url="", n8n_report_url="", n8n_decision_url="") + assert config.n8n_alert_url == "" + + +def test_enforce_env_override(_clean_env: None) -> None: + """DUSK_ENFORCE overrides the default, via the same mechanism as every other field.""" + os.environ["DUSK_ENFORCE"] = "true" + config = load_config("nonexistent-dusk.yaml") + assert config.enforce is True + + +def test_sie_endpoint_env_override(_clean_env: None) -> None: + """DUSK_SIE_ENDPOINT overrides the self-hosted default.""" + os.environ["DUSK_SIE_ENDPOINT"] = "https://hosted-sie.example.com" + config = load_config("nonexistent-dusk.yaml") + assert config.sie_endpoint == "https://hosted-sie.example.com" + + +def test_malformed_yaml_raises(tmp_path: Path) -> None: + """A YAML file that is not a mapping raises ConfigError.""" + yaml_path = tmp_path / "dusk.yaml" + yaml_path.write_text("- just\n- a\n- list\n") + with pytest.raises(ConfigError): + load_config(str(yaml_path)) + + +def test_unknown_key_ignored(tmp_path: Path) -> None: + """Unknown YAML keys are ignored rather than fatal.""" + yaml_path = tmp_path / "dusk.yaml" + yaml_path.write_text("sweep_threshold: 7\nnot_a_real_key: 1\n") + config = load_config(str(yaml_path)) + assert config.sweep_threshold == 7 diff --git a/examples/agent-action-monitor/tests/test_n8n_client.py b/examples/agent-action-monitor/tests/test_n8n_client.py new file mode 100644 index 000000000..b49c2a1e1 --- /dev/null +++ b/examples/agent-action-monitor/tests/test_n8n_client.py @@ -0,0 +1,196 @@ +"""Tests for the n8n webhook client (dusk.trace.n8n_client).""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from dusk.config import Config +from dusk.trace import n8n_client + +_real_get_executor = n8n_client._get_executor + + +class _ImmediateExecutor: + """Runs submitted work synchronously instead of on the real pool, for deterministic tests.""" + + def submit(self, fn, /, *args, **kwargs): # noqa: ANN001, ANN002, ANN003, ANN201 + fn(*args, **kwargs) + + +@pytest.fixture(autouse=True) +def _synchronous_executor(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(n8n_client, "_get_executor", lambda: _ImmediateExecutor()) + + +def _track_send_calls( + monkeypatch: pytest.MonkeyPatch, +) -> list[tuple[str, str, dict[str, object]]]: + calls: list[tuple[str, str, dict[str, object]]] = [] + + def _fake_send(url: str, label: str, payload: dict[str, object]) -> None: + calls.append((url, label, payload)) + + monkeypatch.setattr(n8n_client, "_send", _fake_send) + return calls + + +def test_fire_decision_reads_url_from_config(monkeypatch: pytest.MonkeyPatch) -> None: + calls = _track_send_calls(monkeypatch) + config = Config(n8n_decision_url="https://example.com/decision") + n8n_client.fire_decision({"a": 1}, config=config) + assert calls == [("https://example.com/decision", "decision", {"a": 1})] + + +def test_fire_report_reads_url_from_config(monkeypatch: pytest.MonkeyPatch) -> None: + calls = _track_send_calls(monkeypatch) + config = Config(n8n_report_url="https://example.com/report") + n8n_client.fire_report({"a": 1}, config=config) + assert calls == [("https://example.com/report", "report", {"a": 1})] + + +def test_fire_alert_reads_url_from_config(monkeypatch: pytest.MonkeyPatch) -> None: + calls = _track_send_calls(monkeypatch) + config = Config(n8n_alert_url="https://example.com/alert") + n8n_client.fire_alert({"a": 1}, config=config) + assert calls == [("https://example.com/alert", "alert", {"a": 1})] + + +def test_send_no_op_when_url_empty() -> None: + with patch("urllib.request.urlopen") as mock_urlopen: + n8n_client._send("", "decision", {"a": 1}) + mock_urlopen.assert_not_called() + + +def test_send_rejects_unsupported_scheme() -> None: + with patch("urllib.request.urlopen") as mock_urlopen: + n8n_client._send("ftp://example.com/hook", "decision", {"a": 1}) + mock_urlopen.assert_not_called() + + +def test_send_posts_to_configured_url() -> None: + mock_response = MagicMock() + mock_response.status = 200 + mock_context = MagicMock() + mock_context.__enter__.return_value = mock_response + with patch("urllib.request.urlopen", return_value=mock_context) as mock_urlopen: + n8n_client._send("https://example.com/hook", "decision", {"a": 1}) + mock_urlopen.assert_called_once() + + +def test_send_swallows_errors() -> None: + with patch("urllib.request.urlopen", side_effect=RuntimeError("boom")): + n8n_client._send("https://example.com/hook", "decision", {"a": 1}) + + +def test_webhook_dropped_when_backlog_at_cap( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """Once n8n_max_queued sends are already queued or in flight, a new one is + dropped and logged rather than growing the backlog without limit.""" + calls = _track_send_calls(monkeypatch) + config = Config(n8n_decision_url="https://example.com/decision", n8n_max_queued=2) + monkeypatch.setattr(n8n_client, "_pending_count", 2) + + with caplog.at_level("WARNING"): + n8n_client.fire_decision({"a": 1}, config=config) + + assert calls == [] + assert any("dropped" in r.message for r in caplog.records) + + +def test_pending_count_returns_to_zero_after_delivery(monkeypatch: pytest.MonkeyPatch) -> None: + """The backlog counter must not leak -- a delivered webhook frees its slot.""" + _track_send_calls(monkeypatch) + config = Config(n8n_decision_url="https://example.com/decision") + monkeypatch.setattr(n8n_client, "_pending_count", 0) + + n8n_client.fire_decision({"a": 1}, config=config) + + assert n8n_client._pending_count == 0 + + +def test_fire_with_empty_url_does_not_touch_the_backlog(monkeypatch: pytest.MonkeyPatch) -> None: + calls = _track_send_calls(monkeypatch) + config = Config(n8n_decision_url="") + monkeypatch.setattr(n8n_client, "_pending_count", 0) + + n8n_client.fire_decision({"a": 1}, config=config) + + assert calls == [] + assert n8n_client._pending_count == 0 + + +def test_webhook_backlog_bound_survives_a_real_burst(monkeypatch: pytest.MonkeyPatch) -> None: + """End-to-end proof against the real pool: firing far more webhooks than + n8n_max_queued allows must not grow the backlog past the configured cap. + + Needs the real pool, not the autouse synchronous-executor fixture -- that + fixture runs submitted work inline on the calling thread, which would + make this test's own blocking send deadlock waiting on a release signal + only the same (blocked) thread could ever send. + """ + import threading + import time + + monkeypatch.setattr(n8n_client, "_get_executor", _real_get_executor) + monkeypatch.setattr(n8n_client, "_executor", None) + monkeypatch.setattr(n8n_client, "_pending_count", 0) + config = Config( + n8n_decision_url="https://example.com/decision", n8n_max_workers=1, n8n_max_queued=5 + ) + + release = threading.Event() + max_seen = 0 + seen_lock = threading.Lock() + + def _blocking_send(url: str, label: str, payload: dict[str, object]) -> None: + nonlocal max_seen + with seen_lock: + max_seen = max(max_seen, n8n_client._pending_count) + release.wait(timeout=5) + + monkeypatch.setattr(n8n_client, "_send", _blocking_send) + + for i in range(50): + n8n_client.fire_decision({"i": i}, config=config) + time.sleep(0.1) # let the single worker pick up the first send and start blocking + + with n8n_client._pending_lock: + assert n8n_client._pending_count <= config.n8n_max_queued + + release.set() + + +def test_webhook_concurrency_is_bounded(monkeypatch: pytest.MonkeyPatch) -> None: + """Firing many webhooks in a burst must not spawn one OS thread per call.""" + import threading + import time + + # This test needs the real pool, not the autouse synchronous-executor fixture. + monkeypatch.setattr(n8n_client, "_executor", None) + monkeypatch.setattr(n8n_client, "get_config", lambda: Config(n8n_max_workers=3)) + + in_flight = 0 + max_in_flight = 0 + lock = threading.Lock() + + def _slow_send(url: str, label: str, payload: dict[str, object]) -> None: + nonlocal in_flight, max_in_flight + with lock: + in_flight += 1 + max_in_flight = max(max_in_flight, in_flight) + time.sleep(0.05) + with lock: + in_flight -= 1 + + executor = _real_get_executor() + futures = [ + executor.submit(_slow_send, "https://example.com/decision", "decision", {"i": i}) + for i in range(20) + ] + for f in futures: + f.result(timeout=5) + + assert max_in_flight <= 3 diff --git a/examples/agent-action-monitor/tests/test_sie_live_benchmark.py b/examples/agent-action-monitor/tests/test_sie_live_benchmark.py new file mode 100644 index 000000000..27f6fed85 --- /dev/null +++ b/examples/agent-action-monitor/tests/test_sie_live_benchmark.py @@ -0,0 +1,86 @@ +"""Live SIE benchmark: skipped until DUSK_SIE_ENDPOINT reaches an SIE cluster. + +Once SIE is actually reachable (self-hosted container or the Superlinked-hosted +tester endpoint), this proves the primitives are load-bearing rather than a +no-op: precision/recall on the labelled fixture stay at least as good as the +deterministic-only baseline, and at least one refused action's reasons show a +SIE-sourced signal (rerank or extract), not just the static rule-based checks. +""" + +from __future__ import annotations + +import os +import sys + +import pytest + +LAB_DIR = os.path.join(os.path.dirname(__file__), "..", "lab", "actions") +sys.path.insert(0, os.path.abspath(LAB_DIR)) + +import generate_actions # noqa: E402 + +from dusk.actions.event import AgentAction # noqa: E402 +from dusk.actions.normaliser import normalise_record # noqa: E402 +from dusk.actions.verdict import ActionGate # noqa: E402 +from dusk.config import Config # noqa: E402 +from dusk.trace.vector import sie_encode # noqa: E402 + + +def _normal() -> list[AgentAction]: + return [normalise_record("generic", r) for r in generate_actions.normal_actions()] + + +def _attacks() -> list[AgentAction]: + return [normalise_record("generic", r) for r in generate_actions.out_of_pattern_actions()] + + +@pytest.fixture(autouse=True) +def _skip_unless_sie_reachable() -> None: + if sie_encode("connectivity check") is None: + pytest.skip( + "SIE not installed/reachable; set DUSK_SIE_ENDPOINT and, when required, " + "SIE_API_KEY to run this" + ) + + +def test_live_sie_precision_recall_matches_or_beats_deterministic_baseline() -> None: + config = Config() + gate = ActionGate(config=config) + gate.learn(_normal()) + + labelled = [(a, False) for a in _normal()] + [(a, True) for a in _attacks()] + tp = fp = fn = tn = 0 + for action, is_attack in labelled: + refused = gate.evaluate(action).refused + if is_attack and refused: + tp += 1 + elif is_attack and not refused: + fn += 1 + elif not is_attack and refused: + fp += 1 + else: + tn += 1 + + precision = tp / (tp + fp) if (tp + fp) else 1.0 + recall = tp / (tp + fn) if (tp + fn) else 1.0 + + assert (precision, recall) == (1.0, 1.0), ( + f"precision={precision:.2f} recall={recall:.2f} (tp={tp} fp={fp} fn={fn} tn={tn}) " + "with live SIE -- must not regress versus the deterministic-only baseline" + ) + + +def test_live_sie_primitives_actually_fire_for_at_least_one_attack() -> None: + gate = ActionGate(config=Config()) + gate.learn(_normal()) + + fired_markers = ("SIE rerank", "SIE extract") + saw_sie_signal = any( + any(marker in reason for reason in gate.evaluate(attack).analysis.reasons) + for marker in fired_markers + for attack in _attacks() + ) + assert saw_sie_signal, ( + "expected at least one attack's reasons to carry a SIE rerank/extract marker; " + "if none do, SIE may be configured but not actually contributing a signal" + ) diff --git a/examples/agent-action-monitor/tests/test_trace_vector.py b/examples/agent-action-monitor/tests/test_trace_vector.py new file mode 100644 index 000000000..c525135e7 --- /dev/null +++ b/examples/agent-action-monitor/tests/test_trace_vector.py @@ -0,0 +1,271 @@ +"""Tests for SIE-backed similarity search in dusk.trace.vector.""" + +from __future__ import annotations + +import sys +import types +from unittest.mock import MagicMock + +import pytest + +from dusk.config import Config +from dusk.trace import vector +from dusk.trace.models import TraceDecision + +DEFAULT_CONFIG = Config() + + +def _decisions() -> list[TraceDecision]: + return [ + TraceDecision( + agent_id="netops-agent", + action="firewall_rule_change fw-corp-https", + score=10, + reasoning="opened port 443 on the corp https rule", + ), + TraceDecision( + agent_id="netops-agent", + action="role_assignment fw-restricted", + score=90, + reasoning="granted owner role on a restricted segment", + ), + ] + + +def _inject_fake_item_type(monkeypatch) -> None: + fake_types = types.ModuleType("sie_sdk.types") + fake_types.Item = lambda text=None, id=None, **_kw: { # type: ignore[attr-defined] # noqa: A002 + "text": text, + "id": id, + } + monkeypatch.setitem(sys.modules, "sie_sdk.types", fake_types) + + +def test_find_similar_returns_empty_below_two_decisions() -> None: + assert vector.find_similar("firewall_rule_change fw-corp-https", "netops-agent", []) == [] + + +def test_find_similar_falls_back_to_ngram_when_sie_sdk_missing(monkeypatch) -> None: + monkeypatch.setattr(vector, "_sie_client", lambda config: None) + results = vector.find_similar( + "firewall_rule_change fw-corp-https", "netops-agent", _decisions() + ) + assert isinstance(results, list) + for r in results: + assert isinstance(r, vector.SimilarDecision) + + +def test_stable_hash_is_deterministic_across_calls() -> None: + """_stable_hash must not depend on Python's per-process hash randomization. + + A gate restart while running the no-SIE fallback path must not make + previously recorded embeddings incomparable to freshly computed ones. + """ + assert vector._stable_hash("firewall_rule_change") == vector._stable_hash( + "firewall_rule_change" + ) + assert vector._stable_hash("a") != vector._stable_hash("b") + + +def test_stable_hash_does_not_use_builtin_hash() -> None: + """Guards against a regression back to Python's randomized hash().""" + token = "netops-agent" + assert vector._stable_hash(token) != hash(token) % (2**64) + + +def test_ngram_fallback_is_deterministic_across_calls() -> None: + assert vector._ngram_fallback("firewall_rule_change fw-corp-https") == vector._ngram_fallback( + "firewall_rule_change fw-corp-https" + ) + + +def test_sie_encode_uses_sdk_dense_vector_when_available(monkeypatch) -> None: + fake_client = MagicMock() + fake_client.encode.return_value = {"dense": [1.0, 0.0, 0.0]} + monkeypatch.setattr(vector, "_sie_client", lambda config: fake_client) + _inject_fake_item_type(monkeypatch) + + embedding = vector.sie_encode("hello world") + + assert embedding == [1.0, 0.0, 0.0] + fake_client.encode.assert_called_once() + assert fake_client.encode.call_args[0][0] == DEFAULT_CONFIG.sie_encode_model + + +def test_sie_encode_returns_none_and_does_not_raise_on_sdk_error(monkeypatch) -> None: + fake_client = MagicMock() + fake_client.encode.side_effect = RuntimeError("connection refused") + monkeypatch.setattr(vector, "_sie_client", lambda config: fake_client) + _inject_fake_item_type(monkeypatch) + + assert vector.sie_encode("hello world") is None + + +def test_sie_client_returns_none_when_sie_sdk_not_installed(monkeypatch) -> None: + monkeypatch.setitem(sys.modules, "sie_sdk", None) + assert vector._sie_client(DEFAULT_CONFIG) is None + + +def test_sie_client_passes_configured_timeout(monkeypatch) -> None: + """sie_timeout_ms must actually reach the SDK client, not just live in Config.""" + captured: dict[str, object] = {} + + class FakeSIEClient: + def __init__(self, base_url: str, **kwargs: object) -> None: + captured["base_url"] = base_url + captured.update(kwargs) + + fake_sie_sdk = types.ModuleType("sie_sdk") + fake_sie_sdk.SIEClient = FakeSIEClient # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "sie_sdk", fake_sie_sdk) + + config = Config(sie_endpoint="http://sie:8080", sie_timeout_ms=5000) + client = vector._sie_client(config) + + assert isinstance(client, FakeSIEClient) + assert captured["timeout_s"] == 5.0 + assert captured["base_url"] == "http://sie:8080" + + +def test_sie_score_returns_none_without_candidates() -> None: + assert vector.sie_score("query", []) is None + + +def test_sie_score_returns_none_when_sie_sdk_missing(monkeypatch) -> None: + monkeypatch.setattr(vector, "_sie_client", lambda config: None) + assert vector.sie_score("query", ["a", "b"]) is None + + +def test_sie_score_preserves_input_order(monkeypatch) -> None: + fake_client = MagicMock() + # SDK returns entries out of input order (by rank); sie_score must map + # them back by item_id to the same order the candidates were given in. + fake_client.score.return_value = { + "scores": [ + {"item_id": "1", "score": 0.9, "rank": 0}, + {"item_id": "0", "score": 0.2, "rank": 1}, + ] + } + monkeypatch.setattr(vector, "_sie_client", lambda config: fake_client) + _inject_fake_item_type(monkeypatch) + + scores = vector.sie_score("query", ["candidate-a", "candidate-b"]) + + # Raw logits (0.2, 0.9) pass through sigmoid before returning, so the + # order is preserved but the values are calibrated probabilities. + assert scores is not None + assert scores == [pytest.approx(vector._sigmoid(0.2)), pytest.approx(vector._sigmoid(0.9))] + assert scores[0] < scores[1] + fake_client.score.assert_called_once() + assert fake_client.score.call_args[0][0] == DEFAULT_CONFIG.sie_score_model + + +def test_sie_score_returns_none_and_does_not_raise_on_sdk_error(monkeypatch) -> None: + fake_client = MagicMock() + fake_client.score.side_effect = RuntimeError("connection refused") + monkeypatch.setattr(vector, "_sie_client", lambda config: fake_client) + _inject_fake_item_type(monkeypatch) + + assert vector.sie_score("query", ["a", "b"]) is None + + +def test_sie_extract_returns_empty_when_sie_sdk_missing(monkeypatch) -> None: + monkeypatch.setattr(vector, "_sie_client", lambda config: None) + assert vector.sie_extract("granted owner role") == [] + + +def test_sie_extract_returns_entity_texts_when_available(monkeypatch) -> None: + fake_client = MagicMock() + fake_client.extract.return_value = { + "entities": [ + {"text": "administrator", "label": "role", "score": 0.9}, + {"text": "0.0.0.0", "label": "resource", "score": 0.8}, + ] + } + monkeypatch.setattr(vector, "_sie_client", lambda config: fake_client) + _inject_fake_item_type(monkeypatch) + + terms = vector.sie_extract("grant administrator on 0.0.0.0") + + assert [t.text for t in terms] == ["administrator", "0.0.0.0"] + assert [t.label for t in terms] == ["role", "resource"] + assert [t.score for t in terms] == [0.9, 0.8] + fake_client.extract.assert_called_once() + assert fake_client.extract.call_args[0][0] == DEFAULT_CONFIG.sie_extract_model + assert fake_client.extract.call_args[1]["labels"] == vector.DEFAULT_EXTRACT_LABELS + + +def test_sie_extract_returns_empty_and_does_not_raise_on_sdk_error(monkeypatch) -> None: + fake_client = MagicMock() + fake_client.extract.side_effect = RuntimeError("connection refused") + monkeypatch.setattr(vector, "_sie_client", lambda config: fake_client) + _inject_fake_item_type(monkeypatch) + + assert vector.sie_extract("grant administrator") == [] + + +def test_find_similar_uses_sie_encode_when_available(monkeypatch) -> None: + calls: list[str] = [] + + def fake_encode(text: str, config: Config | None = None) -> list[float]: + calls.append(text) + return [1.0, 0.0] if "fw-corp-https" in text else [0.0, 1.0] + + monkeypatch.setattr(vector, "sie_encode", fake_encode) + results = vector.find_similar( + "firewall_rule_change fw-corp-https", "netops-agent", _decisions() + ) + assert calls + assert all(isinstance(r, vector.SimilarDecision) for r in results) + + +def test_find_similar_reranks_shortlist_with_sie_score(monkeypatch) -> None: + """The rerank pass can override the cosine-similarity order of the shortlist.""" + decisions = [ + TraceDecision(agent_id="netops-agent", action="a", score=10, reasoning="r"), + TraceDecision(agent_id="netops-agent", action="b", score=20, reasoning="r"), + TraceDecision(agent_id="netops-agent", action="c", score=30, reasoning="r"), + ] + monkeypatch.setattr(vector, "sie_encode", lambda text, config=None: [1.0, 0.0]) + + def fake_score(query: str, candidates: list[str]) -> list[float]: + # Same order as candidates: force the last one to the front. + return [0.1, 0.2, 0.9][: len(candidates)] + + monkeypatch.setattr(vector, "sie_score", fake_score) + + results = vector.find_similar("query-action", "netops-agent", decisions, top_k=3) + + assert [r.action for r in results] == ["c", "b", "a"] + + +def test_similar_decision_uses_the_recorded_verdict_not_a_score_guess() -> None: + """Regression test: SimilarDecision.verdict must come from TraceDecision.verdict, + not be reconstructed from a hardcoded score cutoff decoupled from + gate_block_threshold and collapsing WOULD-BLOCK/BLOCK into one label.""" + decisions = [ + TraceDecision( + agent_id="netops-agent", action="a", score=95, reasoning="r", verdict="WOULD-BLOCK" + ), + TraceDecision( + agent_id="netops-agent", action="b", score=10, reasoning="r", verdict="ALLOW" + ), + ] + + scored = list(zip([0.9, 0.8], decisions, strict=True)) + results = vector._rank_candidates("query", scored, top_k=2) + + verdict_by_action = {r.action: r.verdict for r in results} + assert verdict_by_action["a"] == "WOULD-BLOCK" + assert verdict_by_action["b"] == "ALLOW" + + +def test_similar_decision_falls_back_for_legacy_decision_with_no_verdict() -> None: + """A TraceDecision recorded before the verdict field existed has verdict=='' -- + must fall back to a labeled default, not silently claim ALLOW.""" + decisions = [TraceDecision(agent_id="netops-agent", action="a", score=50, reasoning="r")] + scored = [(0.9, decisions[0])] + + results = vector._rank_candidates("query", scored, top_k=1) + + assert results[0].verdict == vector._UNKNOWN_VERDICT_FALLBACK From e476291ecf86878d0ff0923273e2008f0cce3fb4 Mon Sep 17 00:00:00 2001 From: ritiksah141 Date: Mon, 27 Jul 2026 14:27:29 +0100 Subject: [PATCH 2/4] fix: address review feedback on the agent-action-monitor example Rebuilds the live SIE benchmark so it actually shows SIE is load bearing: adds held out negatives the baseline never trained on, and an attack built to evade every deterministic check, then scores it with SIE live versus forced off to show the score lands on opposite sides of the block threshold, not just that a reason string mentions SIE. Also fixes: - SIE_API_KEY and AWS credentials are wired through compose.yml instead of being dead config - the real Bedrock model id now uses the region-prefixed inference profile id Bedrock requires for on-demand invocation - documented the actual cold start cost (weight download size, healthcheck window, RSS, Apple Silicon emulation) - the without-Docker flow now installs agent-demo/mock-prod's own requirements and the gate's default port matches everywhere else - swapped the no-op N8N_USER_MANAGEMENT_DISABLED for N8N_DIAGNOSTICS_ENABLED - a refusal driven only by the repeat-offense boost no longer re-records itself with a fresh timestamp and defeats its own decay - OffenseRecord rejects naive timestamps instead of crashing decay later - the gate engine loader catches OSError broadly, not just FileNotFoundError - an in-process circuit breaker bounds how long a half-up SIE can stall a single gate request - agent_id rejects control characters - offense memory bounds bytes, not just record and agent counts Also merges upstream/main to resolve the gallery table conflict in examples/README.md. --- examples/agent-action-monitor/.env.example | 12 ++- examples/agent-action-monitor/README.md | 83 +++++++++++++------ .../agent-demo/bedrock_client.py | 23 +++-- examples/agent-action-monitor/compose.yml | 22 ++++- .../docs/sie-primitives.md | 36 +++++--- .../lab/actions/generate_actions.py | 56 +++++++++++++ .../src/dusk/actions/analyse.py | 10 +++ .../src/dusk/actions/event.py | 10 ++- .../src/dusk/actions/offense_memory.py | 62 +++++++++++++- .../src/dusk/actions/verdict.py | 13 ++- examples/agent-action-monitor/src/dusk/api.py | 8 +- .../src/dusk/trace/vector.py | 64 ++++++++++++++ .../agent-action-monitor/tests/conftest.py | 15 ++++ .../tests/test_sie_live_benchmark.py | 79 ++++++++++++++++-- 14 files changed, 433 insertions(+), 60 deletions(-) diff --git a/examples/agent-action-monitor/.env.example b/examples/agent-action-monitor/.env.example index cd535e3fb..b8f55fd8b 100644 --- a/examples/agent-action-monitor/.env.example +++ b/examples/agent-action-monitor/.env.example @@ -1,5 +1,7 @@ # Copy to .env and fill in values only when overriding the local defaults. -# Never commit .env. +# Never commit .env. docker compose reads this file automatically for the +# ${VAR} interpolation in compose.yml -- values here reach the containers +# without editing compose.yml itself. # SIE is self-hosted by default and needs no key. Set SIE_API_KEY only for # an authenticated hosted endpoint. @@ -8,3 +10,11 @@ SIE_API_KEY= DUSK_SIE_ENCODE_MODEL=BAAI/bge-m3 DUSK_SIE_SCORE_MODEL=BAAI/bge-reranker-v2-m3 DUSK_SIE_EXTRACT_MODEL=urchade/gliner_multi-v2.1 + +# agent-demo defaults to a mock Bedrock client and needs none of this. Set +# USE_REAL_BEDROCK=true on the agent-demo service in compose.yml (or an +# override file) and fill these in to call real Bedrock instead. +AWS_ACCESS_KEY_ID= +AWS_SECRET_ACCESS_KEY= +AWS_SESSION_TOKEN= +AWS_REGION=us-east-1 diff --git a/examples/agent-action-monitor/README.md b/examples/agent-action-monitor/README.md index ba435eb8a..27803ab52 100644 --- a/examples/agent-action-monitor/README.md +++ b/examples/agent-action-monitor/README.md @@ -45,25 +45,47 @@ target (`mock-prod`), and the agent harness (`agent-demo`) -- all on one internal network, no external egress beyond `sie`'s one-time model-weight download, no API keys required. -The first `up` cold-starts up to three CPU models in `sie` (encode, score, -extract) on demand, not necessarily all at once -- each gate request only -provisions the primitives it actually calls. Every SIE call is bounded to a -short provisioning timeout (1.5s): a cold or at-capacity model falls back -to the deterministic path in about a second rather than blocking the gate, -so a slow or memory-constrained first `sie` startup degrades individual -request latency, it doesn't hang them. Allocate at least 8 GB to Docker -Desktop for `sie` to load its models promptly in the background; under -that, model loading itself takes longer (competing for memory), but -`/v1/gate` keeps responding throughout. +The first `up` downloads roughly 5.7 GB of model weights into the +`sie-hf-cache` volume (the three models in "Model lineup" below), and the +warmed-up `sie` container itself runs at around 7 GB RSS. That download and +initial load has to finish inside the `sie` image's own baked-in healthcheck +window (about 150s of retries) before Docker marks it healthy and lets +`dusk-gate` and `n8n` start (`depends_on: condition: service_healthy` in +compose.yml) -- on a slow connection, or with too little memory allocated to +Docker Desktop, that first `up` can time out with `sie` stuck `unhealthy` +rather than just running slow. Re-running `docker compose up` picks up the +partially-downloaded cache and finishes faster. On Apple Silicon, `sie`'s +image is amd64-only and runs emulated (`platform: linux/amd64` in +compose.yml), which adds meaningfully to both the download-to-ready time and +steady-state latency versus a native amd64 host. + +Allocate at least 8 GB to Docker Desktop. Once `sie` is healthy, cold +individual models are a separate, much smaller concern: each gate request +only provisions the primitives it actually calls, and every SIE call is +bounded to a short provisioning timeout (1.5s) -- a cold or at-capacity +model falls back to the deterministic path in about a second rather than +blocking the gate, so a model that hasn't been hit yet degrades individual +request latency, it doesn't hang them. Without Docker, run the pieces directly. The base install (`pip install -e .`) works on Python 3.11+; the `sie` extra (real SIE encode/score/extract, rather than the deterministic n-gram fallback) requires **Python 3.12+**, since `sie-sdk` itself does -- see the Dockerfile, which uses `python:3.12-slim` for -exactly this reason: +exactly this reason. `agent-demo/` and `mock-prod/` are separate install +targets with their own `requirements.txt` (that's what their own Dockerfiles +install, on top of the base `dusk` package) -- the base install alone is not +enough to run `agent-demo/run_scenario.py`, which needs `requests` (and +`boto3` if you flip on real Bedrock): ```bash -# optional: start from the documented SIE settings +# base package (the gate) +pip install -e . + +# agent-demo and mock-prod are separate install targets, each with their +# own requirements.txt -- the base install above doesn't cover either +pip install -r agent-demo/requirements.txt -r mock-prod/requirements.txt + +# optional: start from the documented SIE/Bedrock settings cp .env.example .env # terminal 1: the gate @@ -122,10 +144,14 @@ to turn it on. Every verdict fires `decision` and `report`; refused verdicts also fire `alert` (`src/dusk/trace/n8n_client.py`). The `n8n` container has these -three webhooks active from startup (`n8n/dusk-webhooks.json`, baked into -the image, not imported by hand) -- each just responds immediately, no -external service involved. Watch them land in the executions list at -`http://localhost:5678`. +three webhooks active from startup (`n8n/dusk-webhooks.json`, imported and +published via the n8n CLI in `n8n/docker-entrypoint.sh`, not imported by +hand, and unaffected by the UI's own owner-account setup since the CLI +talks to n8n's database directly) -- each just responds immediately, no +external service involved. Visiting `http://localhost:5678` in a browser +prompts for a one-time local owner account (n8n's own UI, not something +this stack sets up); the executions list showing the three webhooks firing +is on the other side of that. ## Sample data @@ -183,12 +209,17 @@ record time, capped at 200 entries so lookup cost stays O(1) regardless of how long the gate has been running -- see `src/dusk/api.py`), not hardcoded. This has also been validated against Superlinked's hosted SIE cluster directly, not just assumed: `sie_encode` returns a genuine -1024-dimension `bge-m3` vector, precision/recall on the labelled fixture is -unchanged with SIE live versus the deterministic-only baseline (1.0/1.0 -either way), and at least one attack's reasons carry a real SIE-sourced -marker confirming the primitives are actually contributing a signal over -the network. See `docs/sie-primitives.md` for exactly where each primitive -is wired in. +1024-dimension `bge-m3` vector, and one attack in the live benchmark +(`generate_actions.sie_only_attacks`) is deliberately constructed to evade +every deterministic check -- its deterministic-only score lands under +`gate_block_threshold` by design, so only a SIE signal can refuse it. +Scoring it with real SIE calls versus with them forced off shows the score +lands on opposite sides of the block threshold, not just that a reason +string happens to mention SIE: disabling SIE visibly drops recall on that +attack from 1.0 to 0.0, while precision across the full fixture (including +held-out negatives never folded into the baseline) is unaffected either +way. See `docs/sie-primitives.md` for exactly where each primitive is wired +in and the full benchmark methodology. ## Why SIE specifically @@ -279,9 +310,11 @@ complete local flow: dependent on any AI model at runtime. - SIE's rerank pass only reorders a small shortlist of candidates already retrieved by cosine similarity, not the full decision history. -- The extract model's privileged-term detection is zero-shot and has only - been evaluated against the same synthetic fixtures used elsewhere, not an - adversarial corpus designed to evade it specifically. +- The extract model's privileged-term detection is zero-shot. It has been + shown to catch one deliberately evasive term ("superuser", scored with + and without SIE in `test_disabling_sie_lets_the_evasive_attack_slip_through`), + not evaluated against a broader adversarial corpus designed to evade it + systematically. - Latency numbers are from a single 20-request-per-level trial against a shared tester cluster; enough to confirm the shape, not a high-confidence p95 at every level. See `docs/gate-latency-notes.md`. diff --git a/examples/agent-action-monitor/agent-demo/bedrock_client.py b/examples/agent-action-monitor/agent-demo/bedrock_client.py index e7b5121eb..288f59638 100644 --- a/examples/agent-action-monitor/agent-demo/bedrock_client.py +++ b/examples/agent-action-monitor/agent-demo/bedrock_client.py @@ -2,6 +2,7 @@ from __future__ import annotations +import os from dataclasses import dataclass from typing import Any, Protocol @@ -35,7 +36,11 @@ class DuskBedrockClient: """Wrap a Bedrock-compatible client behind one Converse interface.""" client: BedrockConverseClient - model_id: str = "anthropic.claude-3-5-sonnet-20241022-v2:0" + # Bedrock requires the region-prefixed inference-profile ID for on-demand + # invocation of this model in most regions -- the bare model ID + # ("anthropic.claude-3-5-sonnet-20241022-v2:0") 400s with + # ValidationException: on-demand throughput isn't supported. + model_id: str = "us.anthropic.claude-3-5-sonnet-20241022-v2:0" def converse(self, messages: list[dict[str, Any]]) -> dict[str, Any]: """Call the model and return its raw response. @@ -50,16 +55,22 @@ def converse(self, messages: list[dict[str, Any]]) -> dict[str, Any]: return self.client.converse(modelId=self.model_id, messages=messages) -def build_real_client(region: str = "us-east-1") -> BedrockConverseClient: +def build_real_client(region: str | None = None) -> BedrockConverseClient: """Return a real boto3 bedrock-runtime client. - Requires AWS credentials to be configured in the environment. Only - called when USE_REAL_BEDROCK=true; the default keyless path uses - MockBedrock instead (see mock_bedrock.py, wired in by the harness). + Requires AWS credentials (``AWS_ACCESS_KEY_ID``/``AWS_SECRET_ACCESS_KEY``, + plus ``AWS_SESSION_TOKEN`` for temporary credentials) in the environment + -- compose.yml passes these through from the host's own environment + (empty by default, so the keyless path is unaffected). Only called when + USE_REAL_BEDROCK=true; the default keyless path uses MockBedrock instead + (see mock_bedrock.py, wired in by the harness). Args: - region: AWS region for the client. + region: AWS region for the client. Defaults to the ``AWS_REGION`` + environment variable, then ``"us-east-1"``, matching the ``us.`` + inference-profile prefix on :data:`DuskBedrockClient.model_id`. """ import boto3 + region = region or os.getenv("AWS_REGION", "us-east-1") return boto3.client("bedrock-runtime", region_name=region) # type: ignore[no-any-return] diff --git a/examples/agent-action-monitor/compose.yml b/examples/agent-action-monitor/compose.yml index fd590cd66..0117215dc 100644 --- a/examples/agent-action-monitor/compose.yml +++ b/examples/agent-action-monitor/compose.yml @@ -2,6 +2,10 @@ # up` runs it keyless. Watch mode is the default; DUSK_ENFORCE switches to # enforce. Not egress-blocked (sie needs one-time model-weight downloads on # cold start); nothing here makes an outbound call while serving a request. +# Docker Compose auto-loads a `.env` file next to this one for the ${VAR} +# interpolation below -- `cp .env.example .env` before `docker compose up` +# if you're overriding anything; the defaults (empty SIE_API_KEY, no AWS +# creds) keep the stack keyless exactly as before. services: dusk-gate: @@ -21,6 +25,9 @@ services: DUSK_N8N_REPORT_URL: "http://n8n:5678/webhook/dusk-report" DUSK_N8N_DECISION_URL: "http://n8n:5678/webhook/dusk-decision" DUSK_OFFENSE_MEMORY_PATH: "/app/state/dusk-offense-memory.json" + # Empty by default (self-hosted sie needs no key); set SIE_API_KEY in + # .env only for an authenticated hosted endpoint. + SIE_API_KEY: "${SIE_API_KEY:-}" volumes: - ./sample-data:/app/sample-data:ro - offense-memory:/app/state @@ -45,7 +52,10 @@ services: context: ./n8n # bakes in the three DUSK webhooks, active from startup -- no manual import environment: N8N_SECURE_COOKIE: "false" - N8N_USER_MANAGEMENT_DISABLED: "true" # keyless: no owner-account setup step blocking startup + # No external call is made while serving a request (see header), but + # n8n's own anonymous diagnostics ping n8n.io by default; turn that off + # so that claim holds for the n8n container too. + N8N_DIAGNOSTICS_ENABLED: "false" networks: [internal] ports: ["5678:5678"] # open the n8n UI to inspect the executions list @@ -66,11 +76,19 @@ services: USE_REAL_BEDROCK: "false" DUSK_GATE_URL: "http://dusk-gate:8000/v1/gate" MOCK_PROD_URL: "http://mock-prod:9000/apply" + # Empty by default (mock-Bedrock path ignores them); required in .env + # only when USE_REAL_BEDROCK=true. AWS_SESSION_TOKEN is for temporary + # credentials (SSO/assumed-role) and is fine left empty otherwise. + AWS_ACCESS_KEY_ID: "${AWS_ACCESS_KEY_ID:-}" + AWS_SECRET_ACCESS_KEY: "${AWS_SECRET_ACCESS_KEY:-}" + AWS_SESSION_TOKEN: "${AWS_SESSION_TOKEN:-}" + AWS_REGION: "${AWS_REGION:-us-east-1}" depends_on: dusk-gate: { condition: service_healthy } mock-prod: { condition: service_healthy } networks: [internal] - # mock-Bedrock by default (keyless); USE_REAL_BEDROCK=true for real Bedrock. + # mock-Bedrock by default (keyless); USE_REAL_BEDROCK=true for real + # Bedrock, needs the AWS_* vars above set in .env (see .env.example). networks: internal: diff --git a/examples/agent-action-monitor/docs/sie-primitives.md b/examples/agent-action-monitor/docs/sie-primitives.md index 2b3478adf..5824f364e 100644 --- a/examples/agent-action-monitor/docs/sie-primitives.md +++ b/examples/agent-action-monitor/docs/sie-primitives.md @@ -61,22 +61,33 @@ start behavior. `tests/test_sie_live_benchmark.py` skips until `DUSK_SIE_ENDPOINT` and, for authenticated deployments, `SIE_API_KEY` point at a reachable cluster. Run -against Superlinked's hosted tester endpoint, both checks pass: +against Superlinked's hosted tester endpoint, all checks pass: - `sie_encode` returns a real 1024-dimension dense vector from `BAAI/bge-m3` (confirming the model actually loaded and served, not just that the endpoint answered). -- Precision and recall on the labelled fixture stay at 1.0/1.0, matching the - deterministic-only baseline exactly -- no regression from enabling SIE. -- At least one attack's `reasons` carries a real `SIE rerank` or - `SIE extract` marker, confirming the primitives are actually contributing - a signal over the network, not a no-op that happens to still pass. +- Precision and recall stay at 1.0/1.0 on a fixture that adds two things the + earlier version of this benchmark didn't have: held-out negatives (legitimate + actions never folded into the baseline, not the training set replayed as + its own test) and an attack (`generate_actions.sie_only_attacks`) + constructed so its deterministic-only score lands under + `gate_block_threshold` -- it reuses a known action type and target class, + and its new tokens/values are deliberately kept outside the static + sensitive frozenset. Nothing but a SIE signal can refuse it. +- `test_disabling_sie_lets_the_evasive_attack_slip_through` scores that one + attack twice: once with real SIE calls, once with `sie_extract`/`sie_score` + forced off. It asserts the without-SIE score lands *below* + `gate_block_threshold` and the with-SIE score lands *at or above* it -- + i.e. disabling SIE visibly drops recall on this attack from 1.0 to 0.0, + not just that a reason string happens to mention SIE. - The full test suite passes unchanged with live SIE enabled, confirming nothing depends on the deterministic fallback path being taken. The exact test count is intentionally omitted because it changes as coverage grows. This is evidence that SIE is load-bearing here ("removing SIE degrades the -result"), not just a claim. +result"), not just a claim -- and specifically, not evidence that rests on a +fixture where every attack was already deterministically catchable, which +was the gap in the version of this benchmark reviewed earlier. ## Known limits @@ -98,8 +109,9 @@ result"), not just a claim. - The live decision history behind `similar_decision_ids` is in-memory and capped at 200 entries per gate process -- a demo-scale audit trail, not a durable store. It resets on restart and is not shared across replicas. -- `sie_extract`'s privileged-term detection is zero-shot: it has not been - evaluated against an adversarial corpus designed to evade GLiNER - specifically, only against the same synthetic fixtures used elsewhere. - The 0.5 confidence floor is a reasonable default, not an empirically - tuned threshold. +- `sie_extract`'s privileged-term detection is zero-shot: it has been shown + to catch one deliberately evasive term ("superuser", via + `generate_actions.sie_only_attacks`), not evaluated against a broader + adversarial corpus designed to evade GLiNER systematically. The 0.5 + confidence floor is a reasonable default, not an empirically tuned + threshold. diff --git a/examples/agent-action-monitor/lab/actions/generate_actions.py b/examples/agent-action-monitor/lab/actions/generate_actions.py index fb5746651..8250e6007 100644 --- a/examples/agent-action-monitor/lab/actions/generate_actions.py +++ b/examples/agent-action-monitor/lab/actions/generate_actions.py @@ -95,6 +95,62 @@ def out_of_pattern_actions() -> list[dict[str, Any]]: ] +def held_out_normal_actions() -> list[dict[str, Any]]: + """Legitimate actions from the same three agents, never folded into the baseline. + + Not part of normal_actions() (what the baseline actually learns from) or + out_of_pattern_actions() (labelled attacks) -- used only by the live SIE + benchmark to check precision against negatives the gate could not have + simply memorised, unlike normal_actions() which the benchmark also + trains the baseline on. + """ + return [ + _action(100, "netops-agent", "firewall_rule_change", "fw-corp-smtp", after={"port": 25}), + _action( + 101, + "segment-agent", + "segment_change", + "seg-corporate", + after={"cidr": "10.0.20.0/24"}, + ), + _action( + 102, + "iam-agent", + "role_assignment", + "ra-netops-contributor", + after={"role": "contributor"}, + ), + ] + + +def sie_only_attacks() -> list[dict[str, Any]]: + """Attacks that evade every deterministic check by construction. + + Unlike out_of_pattern_actions(), each of these reuses an action type and + target class the acting agent already has in its baseline, and its new + tokens/change values are deliberately kept outside the static + ``_SENSITIVE_TOKENS``/``_SENSITIVE_VALUES`` frozenset in + ``actions/analyse.py`` -- the point is that the deterministic score alone + must land below ``gate_block_threshold``, so the only thing that can + still catch it is a SIE signal (here, ``extract`` recognising + "superuser" as a role/privilege term the static list doesn't cover; see + ``test_sie_extract_flags_terms_missed_by_the_static_frozenset`` in + ``tests/test_actions_gate.py`` for the same term used against the offline + mock). ``tests/test_sie_live_benchmark.py`` asserts this directly: scored + twice, once with real SIE calls and once with them forced off, rather + than assumed. + """ + return [ + _action( + 103, + "iam-agent", + "role_assignment", + "ra-netops-superuser", + after={"role": "delegate"}, + ), + ] + + def generate(directory: str = FIXTURE_DIR) -> tuple[str, str]: """Write both fixtures into ``directory`` and return their paths.""" os.makedirs(os.path.abspath(directory), exist_ok=True) diff --git a/examples/agent-action-monitor/src/dusk/actions/analyse.py b/examples/agent-action-monitor/src/dusk/actions/analyse.py index 1cf675575..53a0c2996 100644 --- a/examples/agent-action-monitor/src/dusk/actions/analyse.py +++ b/examples/agent-action-monitor/src/dusk/actions/analyse.py @@ -79,6 +79,13 @@ class AnalysisResult: blast_radius: Coarse impact estimate, ``"low"``, ``"medium"`` or ``"high"``. predicted_next: What an attacker would likely do next. + repeat_offense_contribution: The slice of ``score`` contributed by + the repeat-offense signal alone (0.0 when there was no matching + prior offense). Lets a caller ask "would this still be refused + without its own repeat-offense history?" -- used to decide + whether a refusal should be recorded as a new offense at all, + so a refusal caused only by the repeat boost doesn't re-record + itself with a fresh timestamp and defeat its own decay. """ agent_id: str @@ -90,6 +97,7 @@ class AnalysisResult: mitre_atlas: str = "" blast_radius: str = "low" predicted_next: str = "" + repeat_offense_contribution: float = 0.0 def to_dict(self) -> dict[str, Any]: """Return a JSON-serialisable representation of the result.""" @@ -103,6 +111,7 @@ def to_dict(self) -> dict[str, Any]: "mitre_atlas": self.mitre_atlas, "blast_radius": self.blast_radius, "predicted_next": self.predicted_next, + "repeat_offense_contribution": round(self.repeat_offense_contribution, 4), } @@ -376,6 +385,7 @@ def analyse( mitre_atlas=_ATLAS, blast_radius=blast, predicted_next=_predicted_next(action), + repeat_offense_contribution=repeat_score if repeat_reason else 0.0, ) logger.debug( "analysed agent=%s action_type=%s score=%.2f blast=%s", diff --git a/examples/agent-action-monitor/src/dusk/actions/event.py b/examples/agent-action-monitor/src/dusk/actions/event.py index 18329123b..7067bada7 100644 --- a/examples/agent-action-monitor/src/dusk/actions/event.py +++ b/examples/agent-action-monitor/src/dusk/actions/event.py @@ -61,12 +61,16 @@ def __post_init__(self) -> None: """Validate the event, raising ValueError on invalid input. Raises: - ValueError: If ``agent_id`` or ``target`` is empty, ``timestamp`` - is not timezone-aware, ``action_type`` is not a known verb, or - ``change`` is not a mapping. Values are never silently coerced. + ValueError: If ``agent_id`` or ``target`` is empty or contains + control characters (which could forge lines in a log that + cites it verbatim), ``timestamp`` is not timezone-aware, + ``action_type`` is not a known verb, or ``change`` is not a + mapping. Values are never silently coerced. """ if not isinstance(self.agent_id, str) or not self.agent_id.strip(): raise ValueError("agent_id must be a non-empty string") + if any(ord(c) < 0x20 or ord(c) == 0x7F for c in self.agent_id): + raise ValueError("agent_id must not contain control characters") if not isinstance(self.target, str) or not self.target.strip(): raise ValueError("target must be a non-empty string") if not isinstance(self.timestamp, datetime): diff --git a/examples/agent-action-monitor/src/dusk/actions/offense_memory.py b/examples/agent-action-monitor/src/dusk/actions/offense_memory.py index 0f162482f..71c2b304e 100644 --- a/examples/agent-action-monitor/src/dusk/actions/offense_memory.py +++ b/examples/agent-action-monitor/src/dusk/actions/offense_memory.py @@ -17,6 +17,18 @@ #: Retention limits bound both memory use and the persisted file. _MAX_OFFENSES_PER_AGENT = 50 _MAX_TRACKED_AGENTS = 500 +#: Per-record token bounds. The count caps above bound the number of +#: records, not their size -- an adversarial target with thousands of +#: unique tokens would otherwise let a single record's bytes dwarf the +#: caps' intent. Applied at write time, before a record is ever stored. +_MAX_TOKENS_PER_OFFENSE = 32 +_MAX_TOKEN_LENGTH = 64 +#: Hard ceiling on the total in-memory/persisted payload size. The count +#: caps bound records and agents but not bytes; this backstops both against +#: a target crafted to sit just under the per-record token bounds while +#: still being unusually large, and against rotating agent IDs producing +#: many small agents that individually undercut _MAX_TRACKED_AGENTS. +_MAX_TOTAL_BYTES = 10 * 1024 * 1024 @dataclass(frozen=True) @@ -44,6 +56,17 @@ class OffenseRecord: verdict: str timestamp: datetime + def __post_init__(self) -> None: + """Validate the record, raising ValueError on invalid input. + + Raises: + ValueError: If ``timestamp`` is not timezone-aware. A naive + timestamp would later crash every decay calculation for this + agent (aware minus naive datetimes cannot be subtracted). + """ + if self.timestamp.tzinfo is None or self.timestamp.utcoffset() is None: + raise ValueError("timestamp must be timezone-aware") + def to_dict(self) -> dict[str, Any]: """Return a JSON-safe representation.""" data = asdict(self) @@ -77,6 +100,11 @@ class _AgentOffenses: records: list[OffenseRecord] = field(default_factory=list) +def _record_size(entry: OffenseRecord) -> int: + """Approximate on-disk/in-memory footprint of one record, in bytes.""" + return len(json.dumps(entry.to_dict())) + + class OffenseMemory: """Thread-safe offense store with coalesced, single-writer persistence. @@ -96,6 +124,7 @@ def __init__(self, storage_path: str | None = None) -> None: self._storage_path = Path(storage_path) if storage_path else None self._lock = threading.Lock() self._by_agent: OrderedDict[str, _AgentOffenses] = OrderedDict() + self._total_bytes = 0 if self._storage_path is not None: self._load(self._storage_path) self._executor: ThreadPoolExecutor | None = None @@ -114,10 +143,25 @@ def _touch(self, agent_id: str) -> _AgentOffenses: else: self._by_agent.move_to_end(agent_id) while len(self._by_agent) > _MAX_TRACKED_AGENTS: - evicted_id, _ = self._by_agent.popitem(last=False) + evicted_id, evicted = self._by_agent.popitem(last=False) + self._total_bytes -= sum(_record_size(r) for r in evicted.records) logger.info("Evicted offense memory for agent %s (tracked-agent cap)", evicted_id) return offenses + def _enforce_byte_budget(self) -> None: + """Evict the oldest agents until the total payload is back under budget. + + The count caps (per-agent and tracked-agent) bound how many records + and agents exist, not their size -- an adversarial target with many + unique tokens, or enough distinct agent IDs to each stay just under + the tracked-agent cap, can still inflate total bytes well past what + those caps imply. Lock required. + """ + while self._total_bytes > _MAX_TOTAL_BYTES and self._by_agent: + evicted_id, evicted = self._by_agent.popitem(last=False) + self._total_bytes -= sum(_record_size(r) for r in evicted.records) + logger.warning("Evicted offense memory for agent %s (byte-budget cap)", evicted_id) + def _load(self, storage_path: Path) -> None: if not storage_path.exists(): return @@ -151,6 +195,10 @@ def _load(self, storage_path: Path) -> None: self._by_agent[agent_id] = _AgentOffenses(records=capped) while len(self._by_agent) > _MAX_TRACKED_AGENTS: self._by_agent.popitem(last=False) + self._total_bytes = sum( + _record_size(r) for offenses in self._by_agent.values() for r in offenses.records + ) + self._enforce_byte_budget() logger.info( "Loaded offense memory for %d agent(s) from %s", len(self._by_agent), @@ -232,20 +280,27 @@ def record( verdict: str, ) -> None: """Record a refused verdict for ``agent_id`` and schedule a write.""" + bounded_tokens = tuple( + t[:_MAX_TOKEN_LENGTH] for t in sorted(tokens)[:_MAX_TOKENS_PER_OFFENSE] + ) entry = OffenseRecord( trace_id=trace_id, agent_id=agent_id, action_type=action_type, target_class=target_class, - tokens=tuple(sorted(tokens)), + tokens=bounded_tokens, verdict=verdict, timestamp=datetime.now(UTC), ) with self._lock: offenses = self._touch(agent_id) offenses.records.append(entry) + self._total_bytes += _record_size(entry) if len(offenses.records) > _MAX_OFFENSES_PER_AGENT: - del offenses.records[: len(offenses.records) - _MAX_OFFENSES_PER_AGENT] + overflow = offenses.records[: len(offenses.records) - _MAX_OFFENSES_PER_AGENT] + del offenses.records[: len(overflow)] + self._total_bytes -= sum(_record_size(r) for r in overflow) + self._enforce_byte_budget() self._schedule_persist() def offenses_for(self, agent_id: str) -> list[OffenseRecord]: @@ -258,4 +313,5 @@ def clear(self) -> None: """Wipe all recorded offenses. Test-only / operator reset hook.""" with self._lock: self._by_agent.clear() + self._total_bytes = 0 self._schedule_persist() diff --git a/examples/agent-action-monitor/src/dusk/actions/verdict.py b/examples/agent-action-monitor/src/dusk/actions/verdict.py index f3b790130..aadb781e3 100644 --- a/examples/agent-action-monitor/src/dusk/actions/verdict.py +++ b/examples/agent-action-monitor/src/dusk/actions/verdict.py @@ -119,7 +119,18 @@ def evaluate(self, action: AgentAction) -> GateVerdict: else: verdict = ALLOW gate_verdict = GateVerdict(verdict=verdict, analysis=result) - if gate_verdict.refused and self.offense_memory is not None: + # A refusal caused only by the repeat-offense boost (base score below + # the block threshold on its own) is recorded already, via the prior + # offense it matched -- recording it again would give it a fresh + # timestamp and defeat that prior offense's own decay, letting one + # real offense snowball into an indefinite chain of self-renewing + # refusals under watch mode. + base_score = result.score - result.repeat_offense_contribution + if ( + gate_verdict.refused + and self.offense_memory is not None + and base_score >= self.config.gate_block_threshold + ): features = action_features(action) self.offense_memory.record( trace_id=gate_verdict.trace_id, diff --git a/examples/agent-action-monitor/src/dusk/api.py b/examples/agent-action-monitor/src/dusk/api.py index b8ce43bad..c51b54d43 100644 --- a/examples/agent-action-monitor/src/dusk/api.py +++ b/examples/agent-action-monitor/src/dusk/api.py @@ -49,7 +49,7 @@ def _load_gate_engine() -> ActionGate: try: known_good = ingest_file(baseline_path, baseline_source) gate_engine.learn(known_good) - except (FileNotFoundError, ValueError) as exc: + except (OSError, ValueError) as exc: _baseline_load_error = str(exc) logger.error( "gate baseline could not be loaded from %s: %s -- every agent will read as " @@ -220,7 +220,11 @@ def health() -> object: def run() -> None: - port = int(os.getenv("FLASK_PORT", "5000")) + # 8000 everywhere -- Dockerfile, compose.yml, harness.py's default + # DUSK_GATE_URL, and every README example all assume the gate listens + # on 8000; a bare `python -m dusk.api` outside Docker (where the + # Dockerfile's own ENV FLASK_PORT=8000 doesn't apply) must match that. + port = int(os.getenv("FLASK_PORT", "8000")) host = os.getenv("FLASK_HOST", "127.0.0.1") app.run(host=host, port=port) diff --git a/examples/agent-action-monitor/src/dusk/trace/vector.py b/examples/agent-action-monitor/src/dusk/trace/vector.py index de86fd1d4..89357745d 100644 --- a/examples/agent-action-monitor/src/dusk/trace/vector.py +++ b/examples/agent-action-monitor/src/dusk/trace/vector.py @@ -6,6 +6,8 @@ import logging import math import os +import threading +import time from dataclasses import dataclass from typing import Any @@ -20,6 +22,56 @@ #: Per-call provisioning bound; one gate request can make several sequential SIE calls. _PROVISION_TIMEOUT_S = 1.5 +#: A single gate request can make up to five sequential SIE calls (extract, +#: rerank novelty, encode + rerank for similar_decision_ids, encode to +#: record). provision_timeout_s only bounds a cold/at-capacity model; a +#: *half-up* SIE (reachable, but slow to answer) can still let each call run +#: close to the full client timeout, stalling one request for several times +#: sie_timeout_ms. This in-process circuit breaker stops that multiplying +#: effect: after a few consecutive failures, every SIE call is skipped +#: outright (straight to the deterministic/n-gram fallback) for a cooldown +#: window, instead of paying the client timeout again on every call. +_CIRCUIT_FAILURE_THRESHOLD = 3 +_CIRCUIT_COOLDOWN_S = 30.0 + +_circuit_lock = threading.Lock() +_consecutive_sie_failures = 0 +_circuit_open_until = 0.0 + + +def _circuit_is_open() -> bool: + with _circuit_lock: + return time.monotonic() < _circuit_open_until + + +def _record_sie_success() -> None: + global _consecutive_sie_failures, _circuit_open_until + with _circuit_lock: + _consecutive_sie_failures = 0 + _circuit_open_until = 0.0 + + +def _record_sie_failure() -> None: + global _consecutive_sie_failures, _circuit_open_until + with _circuit_lock: + _consecutive_sie_failures += 1 + if _consecutive_sie_failures >= _CIRCUIT_FAILURE_THRESHOLD: + _circuit_open_until = time.monotonic() + _CIRCUIT_COOLDOWN_S + logger.warning( + "SIE circuit opened after %d consecutive failures; skipping SIE " + "calls for %.0fs", + _consecutive_sie_failures, + _CIRCUIT_COOLDOWN_S, + ) + + +def reset_sie_circuit() -> None: + """Close the circuit and clear its failure count. Test-only hook.""" + global _consecutive_sie_failures, _circuit_open_until + with _circuit_lock: + _consecutive_sie_failures = 0 + _circuit_open_until = 0.0 + def _sigmoid(x: float) -> float: """Bound a raw cross-encoder logit into [0, 1]; monotonic, not a calibrated probability.""" @@ -73,6 +125,8 @@ def _sie_client(config: Config) -> Any | None: # noqa: ANN401 def sie_encode(text: str, config: Config | None = None) -> list[float] | None: """encode: text -> dense vector via SIE. Returns None to trigger the n-gram fallback.""" cfg = config or get_config() + if _circuit_is_open(): + return None client = _sie_client(cfg) if client is None: return None @@ -87,8 +141,10 @@ def sie_encode(text: str, config: Config | None = None) -> list[float] | None: provision_timeout_s=_PROVISION_TIMEOUT_S, ) dense = result["dense"] if isinstance(result, dict) else getattr(result, "dense", None) + _record_sie_success() return [float(v) for v in dense] if dense is not None else None except Exception as exc: # noqa: BLE001 + _record_sie_failure() logger.warning("SIE encode failed: %s", exc) return None @@ -105,6 +161,8 @@ def sie_score( if not candidates: return None cfg = config or get_config() + if _circuit_is_open(): + return None client = _sie_client(cfg) if client is None: return None @@ -121,11 +179,13 @@ def sie_score( provision_timeout_s=_PROVISION_TIMEOUT_S, ) entries = result["scores"] if isinstance(result, dict) else getattr(result, "scores", None) + _record_sie_success() if not entries: return None score_by_id = {str(e["item_id"]): _sigmoid(float(e["score"])) for e in entries} return [score_by_id.get(str(i), 0.0) for i in range(len(candidates))] except Exception as exc: # noqa: BLE001 + _record_sie_failure() logger.warning("SIE score failed: %s", exc) return None @@ -141,6 +201,8 @@ def sie_extract( every hit as equally privileged. """ cfg = config or get_config() + if _circuit_is_open(): + return [] client = _sie_client(cfg) if client is None: return [] @@ -158,6 +220,7 @@ def sie_extract( entities = ( result["entities"] if isinstance(result, dict) else getattr(result, "entities", None) ) + _record_sie_success() if not entities: return [] extracted: list[ExtractedTerm] = [] @@ -182,6 +245,7 @@ def sie_extract( ) return extracted except Exception as exc: # noqa: BLE001 + _record_sie_failure() logger.warning("SIE extract failed: %s", exc) return [] diff --git a/examples/agent-action-monitor/tests/conftest.py b/examples/agent-action-monitor/tests/conftest.py index 806d89fc9..a8a700ddd 100644 --- a/examples/agent-action-monitor/tests/conftest.py +++ b/examples/agent-action-monitor/tests/conftest.py @@ -7,6 +7,7 @@ import pytest from dusk.config import reset_config +from dusk.trace.vector import reset_sie_circuit @pytest.fixture(autouse=True) @@ -19,3 +20,17 @@ def _isolate_config() -> Iterator[None]: reset_config() yield reset_config() + + +@pytest.fixture(autouse=True) +def _isolate_sie_circuit() -> Iterator[None]: + """Reset the SIE circuit breaker before and after every test. + + Several tests deliberately make ``sie_encode``/``sie_score``/``sie_extract`` + raise to exercise the fallback path. Without this, three such tests + running back to back in the same process would trip the real circuit + and starve an unrelated, later test of a SIE call it expects to succeed. + """ + reset_sie_circuit() + yield + reset_sie_circuit() diff --git a/examples/agent-action-monitor/tests/test_sie_live_benchmark.py b/examples/agent-action-monitor/tests/test_sie_live_benchmark.py index 27f6fed85..57fb7d7d8 100644 --- a/examples/agent-action-monitor/tests/test_sie_live_benchmark.py +++ b/examples/agent-action-monitor/tests/test_sie_live_benchmark.py @@ -2,15 +2,23 @@ Once SIE is actually reachable (self-hosted container or the Superlinked-hosted tester endpoint), this proves the primitives are load-bearing rather than a -no-op: precision/recall on the labelled fixture stay at least as good as the -deterministic-only baseline, and at least one refused action's reasons show a -SIE-sourced signal (rerank or extract), not just the static rule-based checks. +no-op. Two things the deterministic-only gate cannot claim on its own: + +- Precision holds on held-out negatives (legitimate actions the baseline + never trained on), not just on the training set replayed as its own + negatives, which any gate would trivially get right. +- At least one attack (``generate_actions.sie_only_attacks``) is + constructed to evade every deterministic check, and is scored twice -- + once with real SIE calls, once with them forced off -- to show its score + genuinely lands on opposite sides of the block threshold, not just that a + reason string happens to mention SIE. """ from __future__ import annotations import os import sys +from unittest.mock import patch import pytest @@ -19,6 +27,8 @@ import generate_actions # noqa: E402 +from dusk.actions.analyse import analyse # noqa: E402 +from dusk.actions.baseline import Baseline # noqa: E402 from dusk.actions.event import AgentAction # noqa: E402 from dusk.actions.normaliser import normalise_record # noqa: E402 from dusk.actions.verdict import ActionGate # noqa: E402 @@ -30,10 +40,18 @@ def _normal() -> list[AgentAction]: return [normalise_record("generic", r) for r in generate_actions.normal_actions()] +def _held_out_normal() -> list[AgentAction]: + return [normalise_record("generic", r) for r in generate_actions.held_out_normal_actions()] + + def _attacks() -> list[AgentAction]: return [normalise_record("generic", r) for r in generate_actions.out_of_pattern_actions()] +def _sie_only_attacks() -> list[AgentAction]: + return [normalise_record("generic", r) for r in generate_actions.sie_only_attacks()] + + @pytest.fixture(autouse=True) def _skip_unless_sie_reachable() -> None: if sie_encode("connectivity check") is None: @@ -44,11 +62,25 @@ def _skip_unless_sie_reachable() -> None: def test_live_sie_precision_recall_matches_or_beats_deterministic_baseline() -> None: + """Precision/recall with live SIE, on a fixture the gate could not memorise its way through. + + Negatives include held-out actions never folded into the baseline (not + just the training set replayed as its own test), and attacks include + one constructed to evade every deterministic check (see + generate_actions.sie_only_attacks) -- so a perfect score here reflects + SIE actually contributing, not the fixture being trivial by + construction. + """ config = Config() gate = ActionGate(config=config) gate.learn(_normal()) - labelled = [(a, False) for a in _normal()] + [(a, True) for a in _attacks()] + labelled = ( + [(a, False) for a in _normal()] + + [(a, False) for a in _held_out_normal()] + + [(a, True) for a in _attacks()] + + [(a, True) for a in _sie_only_attacks()] + ) tp = fp = fn = tn = 0 for action, is_attack in labelled: refused = gate.evaluate(action).refused @@ -78,9 +110,46 @@ def test_live_sie_primitives_actually_fire_for_at_least_one_attack() -> None: saw_sie_signal = any( any(marker in reason for reason in gate.evaluate(attack).analysis.reasons) for marker in fired_markers - for attack in _attacks() + for attack in _attacks() + _sie_only_attacks() ) assert saw_sie_signal, ( "expected at least one attack's reasons to carry a SIE rerank/extract marker; " "if none do, SIE may be configured but not actually contributing a signal" ) + + +def test_disabling_sie_lets_the_evasive_attack_slip_through() -> None: + """The property the tests above depend on: SIE is genuinely load-bearing here. + + generate_actions.sie_only_attacks() is built so its deterministic-only + score lands under gate_block_threshold -- it reuses a known action type + and target class, and its new tokens/values are deliberately outside + the static sensitive frozenset. The only thing that can still refuse it + is SIE's extract signal recognising "superuser" as a role/privilege + term. Scoring the same action with real SIE calls versus with them + forced off proves that directly, instead of asserting a reason string + happens to be present. + """ + baseline = Baseline.learn(_normal()) + attack = _sie_only_attacks()[0] + threshold = Config().gate_block_threshold + + with_sie = analyse(baseline, attack) + with ( + patch("dusk.actions.analyse.sie_extract", return_value=[]), + patch("dusk.actions.analyse.sie_score", return_value=None), + ): + without_sie = analyse(baseline, attack) + + assert without_sie.score < threshold, ( + f"expected this attack to evade the deterministic-only score (got " + f"{without_sie.score:.2f} >= {threshold}) -- if it doesn't, it isn't " + f"actually testing what SIE adds; see generate_actions.sie_only_attacks" + ) + assert with_sie.score >= threshold, ( + f"expected live SIE to push this attack's score over the block " + f"threshold (got {with_sie.score:.2f} < {threshold}); if it doesn't, " + f"SIE's extract signal isn't currently catching 'superuser' as a " + f"privileged term -- the load-bearing claim in docs/sie-primitives.md " + f"and the README does not hold for this run" + ) From 0a4a98f9af2ceb6ec441fc6c048147a6880de814 Mon Sep 17 00:00:00 2001 From: ritiksah141 Date: Mon, 27 Jul 2026 15:04:18 +0100 Subject: [PATCH 3/4] fix: address CodeRabbit review comments on the agent-action-monitor example Fixes bugs introduced by the previous commit: - event.py's target field was missing the control-character check its own docstring claimed it had - vector.py recorded a SIE call as successful before its response finished parsing, so a malformed entry double counted as both a success and a failure and skewed the circuit breaker - vulture_whitelist.py referenced a dataclass field with no defaults through the class instead of an instance, which raises AttributeError if the file is ever actually executed Also addresses pre-existing findings surfaced by the same review pass: - run_scenario_or_raise no longer raises for WOULD-BLOCK, since watch mode already applied the action and raising hid that from the caller - DuskBlockedError falls back correctly when reasons is explicitly None, not only when the key is missing - DUSK_ENFORCE is now overridable through .env like the other compose vars - gate_block_threshold is now bounded to <= 1.0, since analyse.py always caps the score to 1.0 and a threshold above it would silently make the gate unable to refuse anything - the n8n webhook executor's lazy singleton now uses the same double-checked locking as the gate engine's, closing a race where two concurrent first calls could each build their own executor - feature extraction now bounds traversal width, not just recursion depth - ingest_file skips any malformed record, not only ones raising AdapterError, matching its own documented one-bad-record contract - the SIE client is now cached per (endpoint, api key, timeout) instead of rebuilt on every single call, and closed when the cache is invalidated - a duplicate assertion and two sleep-based flaky test patterns are cleaned up in favour of explicit synchronization Corrects gate-latency-notes.md, which claimed sie_encode was not on the /v1/gate request path; it is, via _find_similar_decisions/_record_decision. --- .../agent-demo/bedrock_client.py | 4 +- .../agent-demo/harness.py | 11 ++- .../agent-demo/test_harness.py | 15 ++++ examples/agent-action-monitor/compose.yml | 2 +- .../docs/gate-latency-notes.md | 10 ++- .../scripts/vulture_whitelist.py | 6 +- .../src/dusk/actions/baseline.py | 41 +++++++--- .../src/dusk/actions/event.py | 2 + .../src/dusk/actions/ingest.py | 8 +- .../agent-action-monitor/src/dusk/config.py | 10 +++ .../src/dusk/trace/n8n_client.py | 13 ++- .../src/dusk/trace/vector.py | 81 ++++++++++++++++--- .../agent-action-monitor/tests/conftest.py | 17 +++- .../tests/test_actions_gate.py | 1 - .../tests/test_actions_offense_memory.py | 8 +- .../agent-action-monitor/tests/test_config.py | 12 +++ .../tests/test_n8n_client.py | 8 +- 17 files changed, 204 insertions(+), 45 deletions(-) diff --git a/examples/agent-action-monitor/agent-demo/bedrock_client.py b/examples/agent-action-monitor/agent-demo/bedrock_client.py index 288f59638..f278cbb7e 100644 --- a/examples/agent-action-monitor/agent-demo/bedrock_client.py +++ b/examples/agent-action-monitor/agent-demo/bedrock_client.py @@ -16,7 +16,9 @@ class DuskBlockedError(Exception): def __init__(self, verdict: dict[str, Any]) -> None: self.verdict = verdict - reasons = ", ".join(verdict.get("reasons", [])) or "no reason given" + # `or []`, not a `.get` default -- a verdict payload with an explicit + # "reasons": None must fall back the same way a missing key does. + reasons = ", ".join(verdict.get("reasons") or []) or "no reason given" super().__init__(f"blocked ({verdict.get('verdict')}): {reasons}") diff --git a/examples/agent-action-monitor/agent-demo/harness.py b/examples/agent-action-monitor/agent-demo/harness.py index 43fcc2186..434c9d65c 100644 --- a/examples/agent-action-monitor/agent-demo/harness.py +++ b/examples/agent-action-monitor/agent-demo/harness.py @@ -82,12 +82,15 @@ def run_scenario(agent_id: str, scenario: str) -> dict[str, Any]: def run_scenario_or_raise(agent_id: str, scenario: str) -> dict[str, Any]: - """Like run_scenario, but raises DuskBlockedError on WOULD-BLOCK/BLOCK. + """Like run_scenario, but raises DuskBlockedError on a real BLOCK. - Callers that want an exception-based flow (raise on block, proceed on - allow) use this instead of inspecting the returned verdict themselves. + WOULD-BLOCK (watch mode) returns normally instead of raising, since the + action was actually applied -- raising for it would make a caller + believe the action never reached downstream when it did. Callers that + also want to react to WOULD-BLOCK should inspect the returned verdict + themselves. """ result = run_scenario(agent_id, scenario) - if result["verdict"] not in ("ALLOW", "NO_ACTION"): + if result["verdict"] == "BLOCK": raise DuskBlockedError({"verdict": result["verdict"], "reasons": result.get("reasons", [])}) return result diff --git a/examples/agent-action-monitor/agent-demo/test_harness.py b/examples/agent-action-monitor/agent-demo/test_harness.py index 8a1dcb916..e5c0987f4 100644 --- a/examples/agent-action-monitor/agent-demo/test_harness.py +++ b/examples/agent-action-monitor/agent-demo/test_harness.py @@ -96,3 +96,18 @@ def test_run_scenario_or_raise_returns_on_allow(mock_post): result = run_scenario_or_raise("agent-1", "clean") assert result["applied"] is True + + +@patch("harness.requests.post") +def test_run_scenario_or_raise_returns_on_would_block(mock_post): + """WOULD-BLOCK does not raise -- watch mode already applied the action, + so raising would hide that from a caller using the exception-based flow.""" + mock_post.side_effect = [ + _mock_gate_response("WOULD-BLOCK", reasons=["anomalous"]), + _mock_apply_response(), + ] + + result = run_scenario_or_raise("agent-1", "poisoned") + + assert result["verdict"] == "WOULD-BLOCK" + assert result["applied"] is True diff --git a/examples/agent-action-monitor/compose.yml b/examples/agent-action-monitor/compose.yml index 0117215dc..fb15529ef 100644 --- a/examples/agent-action-monitor/compose.yml +++ b/examples/agent-action-monitor/compose.yml @@ -13,7 +13,7 @@ services: context: . dockerfile: Dockerfile environment: - DUSK_ENFORCE: "false" + DUSK_ENFORCE: "${DUSK_ENFORCE:-false}" DUSK_GATE_BASELINE_PATH: "/app/sample-data/baseline.json" DUSK_GATE_BASELINE_SOURCE: "generic" DUSK_DEMO_INTEGRATIONS: "false" diff --git a/examples/agent-action-monitor/docs/gate-latency-notes.md b/examples/agent-action-monitor/docs/gate-latency-notes.md index 23a6acc09..65d950e36 100644 --- a/examples/agent-action-monitor/docs/gate-latency-notes.md +++ b/examples/agent-action-monitor/docs/gate-latency-notes.md @@ -70,10 +70,12 @@ of courtesy to shared, sponsored compute in a visibly degraded state. The hosted tester cluster came back after the outage above, but not into a steady "always warm" state -- it scales its per-model capacity down to zero within roughly a minute of no traffic, then re-provisions on the next -request. `sie_score` and `sie_extract` (the two primitives `/v1/gate` -actually calls per request, via `_extra_sie_signals`; `sie_encode` is not -on this request path) each took 0.1-35s to come back from cold before -settling into sub-second responses. This is a real characteristic of a +request. `sie_score` and `sie_extract` (the two primitives `analyse()` calls +via `_extra_sie_signals`) each took 0.1-35s to come back from cold before +settling into sub-second responses. `sie_encode` is also on the request +path (via `embed_text()` in `_find_similar_decisions`/`_record_decision` in +`api.py`, not through `analyse()`) but was warm throughout this run so it +does not appear in the cold-start figures above. This is a real characteristic of a shared, scale-to-zero tester allocation, not a gate or SDK defect -- `sie_sdk`'s own transient-error retry handled it transparently in every case except when a cold re-provision outlasted `agent-demo/harness.py`'s diff --git a/examples/agent-action-monitor/scripts/vulture_whitelist.py b/examples/agent-action-monitor/scripts/vulture_whitelist.py index 25c91ec8b..09c12c1bb 100644 --- a/examples/agent-action-monitor/scripts/vulture_whitelist.py +++ b/examples/agent-action-monitor/scripts/vulture_whitelist.py @@ -26,4 +26,8 @@ ActionGate.evaluate_all set_config -SimilarDecision.similarity +# SimilarDecision has no default values, so its fields are not class +# attributes -- reference the field through a throwaway instance instead of +# `SimilarDecision.similarity`, which raises AttributeError if this module +# is ever actually executed rather than just parsed by vulture. +SimilarDecision(id="", agent_id="", action="", similarity=0.0, verdict="", score=0).similarity diff --git a/examples/agent-action-monitor/src/dusk/actions/baseline.py b/examples/agent-action-monitor/src/dusk/actions/baseline.py index 66760f4f0..46a3b78d8 100644 --- a/examples/agent-action-monitor/src/dusk/actions/baseline.py +++ b/examples/agent-action-monitor/src/dusk/actions/baseline.py @@ -11,6 +11,14 @@ #: Split a target identifier into lowercase word tokens. _TOKEN_RE = re.compile(r"[a-z0-9]+") +#: Defensive width bounds for request-driven feature extraction. Flask's +#: MAX_CONTENT_LENGTH bounds the request body in bytes, but a pathologically +#: wide or repetitive payload can still sit well under that cap while +#: materializing far more tokens/nodes than any real target or change delta +#: would -- these bound width, on top of the existing recursion-depth guard. +_MAX_TOKENS = 200 +_MAX_FLATTEN_NODES = 500 + def target_class(target: str) -> str: """Return the coarse class of a target, its first token. @@ -24,20 +32,32 @@ class ``"fw"``, while ``"seg-corporate"`` is ``"seg"``. The class captures def target_tokens(target: str) -> set[str]: - """Return the set of lowercase word tokens in a target identifier.""" - return set(_TOKEN_RE.findall(target.lower())) - - -def _flatten_scalars(payload: Any, values: set[str], *, _depth: int = 0) -> None: # noqa: ANN401 - """Collect nested scalar values with a defensive recursion limit.""" - if _depth > 10: + """Return the set of lowercase word tokens in a target identifier, capped at _MAX_TOKENS.""" + return set(_TOKEN_RE.findall(target.lower())[:_MAX_TOKENS]) + + +def _flatten_scalars( + payload: Any, # noqa: ANN401 + values: set[str], + budget: list[int], + *, + _depth: int = 0, +) -> None: + """Collect nested scalar values, bounded by recursion depth and node count. + + ``budget`` is a shared single-element counter (a list so recursive calls + mutate the same cell) decremented once per node visited; traversal stops + once it hits zero, independent of how deep or wide the payload is. + """ + if _depth > 10 or budget[0] <= 0: return + budget[0] -= 1 if isinstance(payload, dict): for value in payload.values(): - _flatten_scalars(value, values, _depth=_depth + 1) + _flatten_scalars(value, values, budget, _depth=_depth + 1) elif isinstance(payload, list): for value in payload: - _flatten_scalars(value, values, _depth=_depth + 1) + _flatten_scalars(value, values, budget, _depth=_depth + 1) elif isinstance(payload, (str, int, float, bool)): values.add(str(payload).lower()) @@ -45,8 +65,9 @@ def _flatten_scalars(payload: Any, values: set[str], *, _depth: int = 0) -> None def _change_values(change: dict[str, Any]) -> set[str]: """Flatten a change delta into a set of stringified scalar values, at any nesting depth.""" values: set[str] = set() + budget = [_MAX_FLATTEN_NODES] for side in ("before", "after"): - _flatten_scalars(change.get(side), values) + _flatten_scalars(change.get(side), values, budget) return values diff --git a/examples/agent-action-monitor/src/dusk/actions/event.py b/examples/agent-action-monitor/src/dusk/actions/event.py index 7067bada7..ed899b9b0 100644 --- a/examples/agent-action-monitor/src/dusk/actions/event.py +++ b/examples/agent-action-monitor/src/dusk/actions/event.py @@ -73,6 +73,8 @@ def __post_init__(self) -> None: raise ValueError("agent_id must not contain control characters") if not isinstance(self.target, str) or not self.target.strip(): raise ValueError("target must be a non-empty string") + if any(ord(c) < 0x20 or ord(c) == 0x7F for c in self.target): + raise ValueError("target must not contain control characters") if not isinstance(self.timestamp, datetime): raise ValueError("timestamp must be a datetime") if self.timestamp.tzinfo is None or self.timestamp.utcoffset() is None: diff --git a/examples/agent-action-monitor/src/dusk/actions/ingest.py b/examples/agent-action-monitor/src/dusk/actions/ingest.py index a28a941ea..2d156b95a 100644 --- a/examples/agent-action-monitor/src/dusk/actions/ingest.py +++ b/examples/agent-action-monitor/src/dusk/actions/ingest.py @@ -29,7 +29,6 @@ def ingest_file(path: str, source: str) -> list[AgentAction]: """ # Imported here so the adapter registry loads lazily and ingest stays free # of an import cycle (ingest -> normaliser -> adapters). - from dusk.actions.adapters.base import AdapterError from dusk.actions.normaliser import normalise_record if not os.path.exists(path): @@ -59,7 +58,12 @@ def ingest_file(path: str, source: str) -> list[AgentAction]: continue try: actions.append(normalise_record(source, record)) - except AdapterError as exc: + except Exception as exc: # noqa: BLE001 + # Deliberately broad: a bug in one adapter for one malformed + # record must not abort the whole batch (this file's own + # contract), and _load_gate_engine() in api.py only catches + # (OSError, ValueError), so anything narrower here would still + # surface as an unhandled 500 through the lazy gate-engine loader. logger.warning("Skipping record %d in %s: %s", index, path, exc) logger.info( diff --git a/examples/agent-action-monitor/src/dusk/config.py b/examples/agent-action-monitor/src/dusk/config.py index 398f76caf..681d9a7fd 100644 --- a/examples/agent-action-monitor/src/dusk/config.py +++ b/examples/agent-action-monitor/src/dusk/config.py @@ -143,6 +143,16 @@ def __post_init__(self) -> None: if value <= 0: raise ConfigError(f"Config value '{name}' must be > 0, got {value!r}") + # analyse.py always caps the anomaly score to [0, 1] (min(1.0, score)), + # so a threshold above 1.0 -- a typo, wrong units, a bad env override + # -- would make the gate permanently unable to refuse anything without + # ever raising a ConfigError to flag it. + if self.gate_block_threshold > 1.0: + raise ConfigError( + f"Config value 'gate_block_threshold' must be <= 1.0, " + f"got {self.gate_block_threshold!r}" + ) + if not isinstance(logging.getLevelName(self.log_level), int): raise ConfigError( f"Config value 'log_level' is not a valid level name: {self.log_level!r}" diff --git a/examples/agent-action-monitor/src/dusk/trace/n8n_client.py b/examples/agent-action-monitor/src/dusk/trace/n8n_client.py index bccdd5a4f..83741ee7f 100644 --- a/examples/agent-action-monitor/src/dusk/trace/n8n_client.py +++ b/examples/agent-action-monitor/src/dusk/trace/n8n_client.py @@ -15,6 +15,7 @@ #: Pool size is fixed from process configuration on first use. _executor: ThreadPoolExecutor | None = None +_executor_lock = threading.Lock() #: Queued-or-in-flight send count; bounds the pool's own unbounded task queue. _pending_lock = threading.Lock() @@ -22,11 +23,17 @@ def _get_executor() -> ThreadPoolExecutor: + # Double-checked locking, matching _get_gate_engine() in api.py: two + # concurrent first calls (e.g. two simultaneous requests at cold start) + # must not each construct their own ThreadPoolExecutor, silently + # discarding one. global _executor if _executor is None: - _executor = ThreadPoolExecutor( - max_workers=get_config().n8n_max_workers, thread_name_prefix="n8n-webhook" - ) + with _executor_lock: + if _executor is None: + _executor = ThreadPoolExecutor( + max_workers=get_config().n8n_max_workers, thread_name_prefix="n8n-webhook" + ) return _executor diff --git a/examples/agent-action-monitor/src/dusk/trace/vector.py b/examples/agent-action-monitor/src/dusk/trace/vector.py index 89357745d..e94180591 100644 --- a/examples/agent-action-monitor/src/dusk/trace/vector.py +++ b/examples/agent-action-monitor/src/dusk/trace/vector.py @@ -73,6 +73,12 @@ def reset_sie_circuit() -> None: _circuit_open_until = 0.0 +#: Cached SIEClient, rebuilt only when (endpoint, api_key, timeout) changes. +_client_lock = threading.Lock() +_cached_sie_client: Any | None = None +_cached_sie_client_key: tuple[str, str | None, int] | None = None + + def _sigmoid(x: float) -> float: """Bound a raw cross-encoder logit into [0, 1]; monotonic, not a calibrated probability.""" try: @@ -105,21 +111,66 @@ class SimilarDecision: def _sie_client(config: Config) -> Any | None: # noqa: ANN401 - """Return a constructed SIEClient if the sie-sdk package is installed, else None.""" + """Return a cached SIEClient for this (endpoint, api_key, timeout) triple. + + A fresh client per call was the original shape here; reused instead, so + a sustained-traffic gate process doesn't construct and leak one client + (and whatever connection it opens) per SIE call. Rebuilt only when the + resolved settings actually change -- e.g. a config reload -- closing the + stale client first via whatever cleanup the SDK exposes. + """ + # Checked on every call, not cached: this import is what + # test_sie_client_returns_none_when_sie_sdk_not_installed simulates + # being absent via sys.modules, and Python's own import cache already + # makes a repeated `from ... import` here effectively free. try: from sie_sdk import SIEClient # type: ignore[import-not-found] except ImportError: return None + + api_key = os.getenv("SIE_API_KEY") or None + key = (config.sie_endpoint, api_key, config.sie_timeout_ms) + + global _cached_sie_client, _cached_sie_client_key + with _client_lock: + if _cached_sie_client is not None and _cached_sie_client_key == key: + return _cached_sie_client + if _cached_sie_client is not None: + _close_sie_client(_cached_sie_client) + _cached_sie_client = None + _cached_sie_client_key = None + try: + client = SIEClient( + config.sie_endpoint.rstrip("/"), + api_key=api_key, + timeout_s=config.sie_timeout_ms / 1000, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("SIE client construction failed: %s", exc) + return None + _cached_sie_client = client + _cached_sie_client_key = key + return client + + +def _close_sie_client(client: Any) -> None: # noqa: ANN401 + close = getattr(client, "close", None) + if not callable(close): + return try: - api_key = os.getenv("SIE_API_KEY") or None - return SIEClient( - config.sie_endpoint.rstrip("/"), - api_key=api_key, - timeout_s=config.sie_timeout_ms / 1000, - ) + close() except Exception as exc: # noqa: BLE001 - logger.warning("SIE client construction failed: %s", exc) - return None + logger.warning("Error closing stale SIE client: %s", exc) + + +def reset_sie_client() -> None: + """Close and drop the cached SIEClient. Test-only hook, also usable on shutdown.""" + global _cached_sie_client, _cached_sie_client_key + with _client_lock: + if _cached_sie_client is not None: + _close_sie_client(_cached_sie_client) + _cached_sie_client = None + _cached_sie_client_key = None def sie_encode(text: str, config: Config | None = None) -> list[float] | None: @@ -141,8 +192,9 @@ def sie_encode(text: str, config: Config | None = None) -> list[float] | None: provision_timeout_s=_PROVISION_TIMEOUT_S, ) dense = result["dense"] if isinstance(result, dict) else getattr(result, "dense", None) + parsed = [float(v) for v in dense] if dense is not None else None _record_sie_success() - return [float(v) for v in dense] if dense is not None else None + return parsed except Exception as exc: # noqa: BLE001 _record_sie_failure() logger.warning("SIE encode failed: %s", exc) @@ -179,11 +231,13 @@ def sie_score( provision_timeout_s=_PROVISION_TIMEOUT_S, ) entries = result["scores"] if isinstance(result, dict) else getattr(result, "scores", None) - _record_sie_success() if not entries: + _record_sie_success() return None score_by_id = {str(e["item_id"]): _sigmoid(float(e["score"])) for e in entries} - return [score_by_id.get(str(i), 0.0) for i in range(len(candidates))] + parsed = [score_by_id.get(str(i), 0.0) for i in range(len(candidates))] + _record_sie_success() + return parsed except Exception as exc: # noqa: BLE001 _record_sie_failure() logger.warning("SIE score failed: %s", exc) @@ -220,8 +274,8 @@ def sie_extract( entities = ( result["entities"] if isinstance(result, dict) else getattr(result, "entities", None) ) - _record_sie_success() if not entities: + _record_sie_success() return [] extracted: list[ExtractedTerm] = [] for e in entities: @@ -243,6 +297,7 @@ def sie_extract( score=float(getattr(e, "score", 0.0)), ) ) + _record_sie_success() return extracted except Exception as exc: # noqa: BLE001 _record_sie_failure() diff --git a/examples/agent-action-monitor/tests/conftest.py b/examples/agent-action-monitor/tests/conftest.py index a8a700ddd..f0dddc96f 100644 --- a/examples/agent-action-monitor/tests/conftest.py +++ b/examples/agent-action-monitor/tests/conftest.py @@ -7,7 +7,7 @@ import pytest from dusk.config import reset_config -from dusk.trace.vector import reset_sie_circuit +from dusk.trace.vector import reset_sie_circuit, reset_sie_client @pytest.fixture(autouse=True) @@ -34,3 +34,18 @@ def _isolate_sie_circuit() -> Iterator[None]: reset_sie_circuit() yield reset_sie_circuit() + + +@pytest.fixture(autouse=True) +def _isolate_sie_client() -> Iterator[None]: + """Reset the cached SIEClient before and after every test. + + ``_sie_client`` now caches its client across calls, keyed on + (endpoint, api_key, timeout). Without this, a test that fakes a client + for one (endpoint, api_key, timeout) triple -- e.g. by monkeypatching + ``sys.modules['sie_sdk']`` -- could bleed its cached fake into a later + test that happens to use the same triple with real sie-sdk installed. + """ + reset_sie_client() + yield + reset_sie_client() diff --git a/examples/agent-action-monitor/tests/test_actions_gate.py b/examples/agent-action-monitor/tests/test_actions_gate.py index eeaf94f91..ce7801401 100644 --- a/examples/agent-action-monitor/tests/test_actions_gate.py +++ b/examples/agent-action-monitor/tests/test_actions_gate.py @@ -265,7 +265,6 @@ def test_sie_extract_high_confidence_non_privileged_label_is_ignored() -> None: assert extracted.score == baseline_only.score assert not any("SIE extract" in r for r in extracted.reasons) - assert not any("SIE extract" in r for r in extracted.reasons) def test_sie_extract_unavailable_does_not_change_score() -> None: diff --git a/examples/agent-action-monitor/tests/test_actions_offense_memory.py b/examples/agent-action-monitor/tests/test_actions_offense_memory.py index ac7f538ee..9e118d621 100644 --- a/examples/agent-action-monitor/tests/test_actions_offense_memory.py +++ b/examples/agent-action-monitor/tests/test_actions_offense_memory.py @@ -149,8 +149,12 @@ def test_record_does_not_block_on_a_slow_disk_write(tmp_path: Path, monkeypatch) storage = tmp_path / "offenses.json" original_replace = Path.replace + # A wide gap between the patched delay and the assertion threshold below + # -- not the delay itself -- is what makes this robust under CI + # contention; a delay only as large as the threshold leaves no margin + # for record()'s own scheduling overhead. def slow_replace(self: Path, target: object) -> object: - time.sleep(0.2) + time.sleep(1.0) return original_replace(self, target) monkeypatch.setattr(Path, "replace", slow_replace) @@ -160,7 +164,7 @@ def slow_replace(self: Path, target: object) -> object: _record(memory, trace_id="fast-return") elapsed = time.monotonic() - start - assert elapsed < 0.2, "record() waited on the disk write instead of backgrounding it" + assert elapsed < 0.5, "record() waited on the disk write instead of backgrounding it" assert not storage.exists(), "write should still be in flight at this point" memory.flush() diff --git a/examples/agent-action-monitor/tests/test_config.py b/examples/agent-action-monitor/tests/test_config.py index a87cf1965..c2d43b7ea 100644 --- a/examples/agent-action-monitor/tests/test_config.py +++ b/examples/agent-action-monitor/tests/test_config.py @@ -82,6 +82,18 @@ def test_invalid_threshold_raises() -> None: Config(sweep_threshold=0) +def test_gate_block_threshold_above_one_raises() -> None: + """A threshold above 1.0 would make the gate never refuse anything, so it's rejected.""" + with pytest.raises(ConfigError): + Config(gate_block_threshold=1.5) + + +def test_gate_block_threshold_of_exactly_one_is_valid() -> None: + """1.0 is the inclusive upper bound, not an off-by-one exclusion.""" + config = Config(gate_block_threshold=1.0) + assert config.gate_block_threshold == 1.0 + + def test_invalid_log_level_raises() -> None: """An unrecognised log level name is rejected.""" with pytest.raises(ConfigError): diff --git a/examples/agent-action-monitor/tests/test_n8n_client.py b/examples/agent-action-monitor/tests/test_n8n_client.py index b49c2a1e1..06cd608bf 100644 --- a/examples/agent-action-monitor/tests/test_n8n_client.py +++ b/examples/agent-action-monitor/tests/test_n8n_client.py @@ -132,7 +132,6 @@ def test_webhook_backlog_bound_survives_a_real_burst(monkeypatch: pytest.MonkeyP only the same (blocked) thread could ever send. """ import threading - import time monkeypatch.setattr(n8n_client, "_get_executor", _real_get_executor) monkeypatch.setattr(n8n_client, "_executor", None) @@ -142,6 +141,7 @@ def test_webhook_backlog_bound_survives_a_real_burst(monkeypatch: pytest.MonkeyP ) release = threading.Event() + entered = threading.Event() max_seen = 0 seen_lock = threading.Lock() @@ -149,13 +149,17 @@ def _blocking_send(url: str, label: str, payload: dict[str, object]) -> None: nonlocal max_seen with seen_lock: max_seen = max(max_seen, n8n_client._pending_count) + entered.set() release.wait(timeout=5) monkeypatch.setattr(n8n_client, "_send", _blocking_send) for i in range(50): n8n_client.fire_decision({"i": i}, config=config) - time.sleep(0.1) # let the single worker pick up the first send and start blocking + # Deterministic instead of a fixed sleep: wait for the single worker to + # actually pick up the first send and enter its blocking wait, rather + # than guessing how long that scheduling takes under CI contention. + assert entered.wait(timeout=5), "worker never picked up the first send" with n8n_client._pending_lock: assert n8n_client._pending_count <= config.n8n_max_queued From 4121770019e755c2bb3fd9c2e4c3f9362096175c Mon Sep 17 00:00:00 2001 From: ritiksah141 Date: Mon, 27 Jul 2026 15:13:20 +0100 Subject: [PATCH 4/4] fix: bound token/node scanning work, not just retained output, in baseline.py target_tokens() used findall()[:_MAX_TOKENS], which still builds the full match list before the slice runs, so a repetitive target could force an unbounded match list even though only the first _MAX_TOKENS ever reached the returned set. Switched to finditer() + islice() so the regex work itself is bounded. _flatten_scalars() checked its budget before decrementing for depth, so a depth-rejected node never consumed budget, and the enclosing dict/list loops kept enumerating every remaining item as a wasted no-op call once the budget hit zero instead of stopping. Now decrements before the depth check and breaks out of the loops as soon as the budget is exhausted. --- .../src/dusk/actions/baseline.py | 28 +++++++++++++++---- 1 file changed, 23 insertions(+), 5 deletions(-) diff --git a/examples/agent-action-monitor/src/dusk/actions/baseline.py b/examples/agent-action-monitor/src/dusk/actions/baseline.py index 46a3b78d8..f925b438c 100644 --- a/examples/agent-action-monitor/src/dusk/actions/baseline.py +++ b/examples/agent-action-monitor/src/dusk/actions/baseline.py @@ -4,6 +4,7 @@ import re from dataclasses import dataclass, field +from itertools import islice from typing import Any from dusk.actions.event import AgentAction @@ -32,8 +33,15 @@ class ``"fw"``, while ``"seg-corporate"`` is ``"seg"``. The class captures def target_tokens(target: str) -> set[str]: - """Return the set of lowercase word tokens in a target identifier, capped at _MAX_TOKENS.""" - return set(_TOKEN_RE.findall(target.lower())[:_MAX_TOKENS]) + """Return the set of lowercase word tokens in a target identifier, capped at _MAX_TOKENS. + + Uses ``finditer`` + ``islice`` rather than ``findall()[:_MAX_TOKENS]`` -- + ``findall`` would materialize every match before the slice ever runs, so + a repetitive target could still force an unbounded match list even + though only the first _MAX_TOKENS ever reach the returned set. + """ + matches = islice(_TOKEN_RE.finditer(target.lower()), _MAX_TOKENS) + return {m.group(0) for m in matches} def _flatten_scalars( @@ -46,17 +54,27 @@ def _flatten_scalars( """Collect nested scalar values, bounded by recursion depth and node count. ``budget`` is a shared single-element counter (a list so recursive calls - mutate the same cell) decremented once per node visited; traversal stops - once it hits zero, independent of how deep or wide the payload is. + mutate the same cell) decremented once per node visited, including + depth-rejected ones -- otherwise a wide container sitting right at the + depth boundary could be fully scanned without ever touching the budget. + Traversal stops once the budget hits zero: both the recursive calls + return immediately, and the loops below break rather than continuing to + enumerate a large dict/list whose items would each be a wasted no-op call. """ - if _depth > 10 or budget[0] <= 0: + if budget[0] <= 0: return budget[0] -= 1 + if _depth > 10: + return if isinstance(payload, dict): for value in payload.values(): + if budget[0] <= 0: + break _flatten_scalars(value, values, budget, _depth=_depth + 1) elif isinstance(payload, list): for value in payload: + if budget[0] <= 0: + break _flatten_scalars(value, values, budget, _depth=_depth + 1) elif isinstance(payload, (str, int, float, bool)): values.add(str(payload).lower())