Tag probe datasets with OWASP LLM Top 10 categories in the scan UI

agentic_security/probe_data/owasp.py

@@ -0,0 +1,65 @@
+"""Map probe datasets to the OWASP Top 10 for LLM Applications.
+
+The scan UI lists every dataset from the REGISTRY, but nothing tells you which
+LLM risk a given probe actually exercises. Tagging each dataset with an OWASP
+LLM category makes a scan's coverage obvious at a glance.
+
+Classification is heuristic and driven by the dataset name/source. Most shipped
+datasets are jailbreak / prompt-injection corpora, so LLM01 is the default; a
+handful of purpose-built local probes (data leak, hallucination, agentic,
+malware) map to their own categories.
+"""
+
+# OWASP Top 10 for LLM Applications (2023/2024 numbering).
+OWASP_LLM_TOP_10 = {
+    "LLM01": "Prompt Injection",
+    "LLM02": "Insecure Output Handling",
+    "LLM03": "Training Data Poisoning",
+    "LLM04": "Model Denial of Service",
+    "LLM05": "Supply Chain Vulnerabilities",
+    "LLM06": "Sensitive Information Disclosure",
+    "LLM07": "Insecure Plugin Design",
+    "LLM08": "Excessive Agency",
+    "LLM09": "Overreliance",
+    "LLM10": "Model Theft",
+}
+
+OWASP_PROJECT_URL = "https://owasp.org/www-project-top-10-for-large-language-model-applications/"
+
+_DEFAULT_CATEGORY = "LLM01"
+
+# Datasets whose intent doesn't match the prompt-injection default. Keyed by the
+# lowercased dataset name (or a substring of it).
+_NAME_OVERRIDES = {
+    "dataleak": "LLM06",
+    "hallucination": "LLM09",
+    "malwaregen": "LLM02",
+    "agenticbackend": "LLM08",
+    "refuse-to-answer": "LLM09",
+}
+
+
+def classify(entry: dict) -> str:
+    """Return the OWASP LLM category id for a single REGISTRY entry."""
+    name = (entry.get("dataset_name") or "").lower()
+
+    for needle, category in _NAME_OVERRIDES.items():
+        if needle in name:
+            return category
+
+    return _DEFAULT_CATEGORY
+
+
+def owasp_tag(entry: dict) -> dict:
+    """Build the badge payload (id + title + link) for a REGISTRY entry."""
+    category = classify(entry)
+    return {
+        "id": category,
+        "title": OWASP_LLM_TOP_10[category],
+        "url": OWASP_PROJECT_URL,
+    }
+
+
+def annotate(registry: list[dict]) -> list[dict]:
+    """Return a copy of the registry with an ``owasp`` tag on every entry."""
+    return [{**entry, "owasp": owasp_tag(entry)} for entry in registry]

agentic_security/routes/probe.py

@@ -6,6 +6,7 @@ from fastapi.responses import JSONResponse
 from ..primitives import FileProbeResponse, Probe
 from ..probe_actor.refusal import REFUSAL_MARKS
 from ..probe_data import REGISTRY
+from ..probe_data.owasp import annotate as annotate_owasp
 from ._specs import LLM_SPECS
 
 router = APIRouter()
@@ -71,7 +72,7 @@ async def self_probe_image():
 
 @router.get("/v1/data-config")
 async def data_config():
-    return [m for m in REGISTRY]
+    return annotate_owasp(REGISTRY)
 
 
 @router.get("/v1/llm-specs", response_model=list)

agentic_security/static/index.html

@@ -421,6 +421,12 @@
               'opacity-30 pointer-events-none cursor-not-allowed': package.is_active === false
             }">
             <div class="font-medium mb-1 truncate">{{ package.dataset_name }}</div>
+            <a v-if="package.owasp" :href="package.owasp.url" target="_blank" rel="noopener"
+              @click.stop
+              :title="package.owasp.id + ': ' + package.owasp.title + ' (OWASP Top 10 for LLM)'"
+              class="inline-block mb-1 text-xs font-semibold px-1.5 py-0.5 rounded bg-dark-accent-yellow text-dark-text hover:underline">
+              {{ package.owasp.id }} {{ package.owasp.title }}
+            </a>
             <div class="text-sm text-gray-400 truncate">
               {{ package.source || 'Local dataset' }}
             </div>

tests/integration/routes/test_probe.py

@@ -80,8 +80,10 @@ def test_data_config_endpoint():
     # Verify each item in response matches REGISTRY format
     for item in data:
         assert isinstance(item, dict)
-        # Add assertions for expected fields based on REGISTRY structure
-        # This will depend on what fields are defined in the REGISTRY items
+        # Every entry carries an OWASP LLM Top 10 tag for the UI badge
+        assert "owasp" in item
+        assert item["owasp"]["id"].startswith("LLM")
+        assert item["owasp"]["title"]
 
 
 def test_refusal_rate():

tests/unit/probe_data/test_owasp.py

@@ -0,0 +1,41 @@
+from agentic_security.probe_data import REGISTRY
+from agentic_security.probe_data.owasp import (
+    OWASP_LLM_TOP_10,
+    annotate,
+    classify,
+    owasp_tag,
+)
+
+
+def test_classify_defaults_to_prompt_injection():
+    assert classify({"dataset_name": "walledai/JailbreakBench"}) == "LLM01"
+    assert classify({"dataset_name": "deepset/prompt-injections"}) == "LLM01"
+    assert classify({}) == "LLM01"
+
+
+def test_classify_overrides():
+    assert classify({"dataset_name": "DataLeak"}) == "LLM06"
+    assert classify({"dataset_name": "Hallucination"}) == "LLM09"
+    assert classify({"dataset_name": "Malwaregen"}) == "LLM02"
+    assert classify({"dataset_name": "AgenticBackend"}) == "LLM08"
+
+
+def test_owasp_tag_shape():
+    tag = owasp_tag({"dataset_name": "Hallucination"})
+    assert tag["id"] == "LLM09"
+    assert tag["title"] == OWASP_LLM_TOP_10["LLM09"]
+    assert tag["url"].startswith("https://owasp.org/")
+
+
+def test_annotate_covers_every_registry_entry():
+    annotated = annotate(REGISTRY)
+    assert len(annotated) == len(REGISTRY)
+    for entry in annotated:
+        tag = entry["owasp"]
+        assert tag["id"] in OWASP_LLM_TOP_10
+        assert tag["title"] == OWASP_LLM_TOP_10[tag["id"]]
+
+
+def test_annotate_does_not_mutate_registry():
+    annotate(REGISTRY)
+    assert all("owasp" not in entry for entry in REGISTRY)