diff --git a/agentic_security/probe_data/owasp.py b/agentic_security/probe_data/owasp.py new file mode 100644 index 0000000..dc695c9 --- /dev/null +++ b/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] diff --git a/agentic_security/routes/probe.py b/agentic_security/routes/probe.py index 3dee380..5962b8a 100644 --- a/agentic_security/routes/probe.py +++ b/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) diff --git a/agentic_security/static/index.html b/agentic_security/static/index.html index 0a332bc..9a7d60a 100644 --- a/agentic_security/static/index.html +++ b/agentic_security/static/index.html @@ -421,6 +421,12 @@ 'opacity-30 pointer-events-none cursor-not-allowed': package.is_active === false }">
{{ package.dataset_name }}
+ + {{ package.owasp.id }} {{ package.owasp.title }} +
{{ package.source || 'Local dataset' }}
diff --git a/tests/integration/routes/test_probe.py b/tests/integration/routes/test_probe.py index fd1fe5b..9fadd3b 100644 --- a/tests/integration/routes/test_probe.py +++ b/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(): diff --git a/tests/unit/probe_data/test_owasp.py b/tests/unit/probe_data/test_owasp.py new file mode 100644 index 0000000..e8a0451 --- /dev/null +++ b/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)