feat: add compliance framework tagging for AI red teaming (#318)

Add OWASP, ATLAS, SAIF, and NIST compliance tags to attacks and transforms.
Attacks tagged with core jailbreak technique (LLM01), transforms tagged with
specific vulnerability categories. Includes comprehensive test coverage.

Author: rdheekonda

dreadnode/airt/__init__.py

@@ -1,4 +1,4 @@
-from dreadnode.airt import attack, search
+from dreadnode.airt import attack, compliance, search
 from dreadnode.airt.attack import (
     Attack,
     goat_attack,
@@ -9,20 +9,37 @@
     tap_attack,
     zoo_attack,
 )
+from dreadnode.airt.compliance import (
+    ATTACK_MAPPINGS,
+    ATLASTechnique,
+    NISTAIRMFFunction,
+    OWASPCategory,
+    SAIFCategory,
+    tag_attack,
+    tag_transform,
+)
 from dreadnode.airt.target import CustomTarget, LLMTarget, Target
 
 __all__ = [
+    "ATTACK_MAPPINGS",
+    "ATLASTechnique",
     "Attack",
     "CustomTarget",
     "LLMTarget",
+    "NISTAIRMFFunction",
+    "OWASPCategory",
+    "SAIFCategory",
     "Target",
     "attack",
+    "compliance",
     "goat_attack",
     "hop_skip_jump_attack",
     "nes_attack",
     "prompt_attack",
     "search",
     "simba_attack",
+    "tag_attack",
+    "tag_transform",
     "tap_attack",
     "target",
     "zoo_attack",

dreadnode/airt/attack/base.py

@@ -26,6 +26,9 @@ class Attack(Study[In, Out]):
     tags: list[str] = Config(default_factory=lambda: ["attack"])
     """A list of tags associated with the attack for logging."""
 
+    compliance_tags: dict[str, t.Any] = Config(default_factory=dict)
+    """Compliance framework tags (OWASP, ATLAS, SAIF, NIST) for this attack."""
+
     hooks: list[EvalHook] = Field(default_factory=list, exclude=True, repr=False)
     """Hooks to run at various points in the attack lifecycle."""
 

dreadnode/airt/attack/crescendo.py

@@ -3,6 +3,13 @@
 import yaml
 
 from dreadnode.airt.attack import Attack
+from dreadnode.airt.compliance import (
+    ATLASTechnique,
+    NISTAIRMFFunction,
+    OWASPCategory,
+    SAIFCategory,
+    tag_attack,
+)
 from dreadnode.airt.target.base import Target
 from dreadnode.constants import CRESCENDO_VARIANT_1
 from dreadnode.data_types.message import Message as DnMessage
@@ -14,6 +21,20 @@
 from dreadnode.transforms.base import Transform
 from dreadnode.transforms.refine import adapt_prompt_trials, llm_refine
 
+# Compliance framework tags for Crescendo attack
+# Core jailbreak technique tags - specific vulnerability categories (LLM02, LLM07, etc.)
+# are added when transforms targeting those categories are used
+COMPLIANCE_TAGS = tag_attack(
+    atlas=[
+        ATLASTechnique.PROMPT_INJECTION_DIRECT,
+        ATLASTechnique.LLM_JAILBREAK,
+    ],
+    owasp=OWASPCategory.LLM01_PROMPT_INJECTION,
+    saif=SAIFCategory.INPUT_MANIPULATION,
+    nist_function=NISTAIRMFFunction.MEASURE,
+    nist_subcategory="MS-2.7",
+)
+
 
 def crescendo_attack(
     goal: str,
@@ -179,6 +200,7 @@ async def crescendo_refiner(trials: list[Trial[DnMessage]]) -> DnMessage:
             "objective": objective_judge,
         },
         hooks=hooks or [],
+        compliance_tags=COMPLIANCE_TAGS,
     )
 
     # Add stop condition based on early_stopping_score

dreadnode/airt/attack/goat.py

@@ -1,6 +1,13 @@
 import typing as t
 
 from dreadnode.airt.attack import Attack
+from dreadnode.airt.compliance import (
+    ATLASTechnique,
+    NISTAIRMFFunction,
+    OWASPCategory,
+    SAIFCategory,
+    tag_attack,
+)
 from dreadnode.data_types.message import Message as DnMessage
 from dreadnode.meta.context import TrialCandidate
 from dreadnode.optimization.search.graph import graph_neighborhood_search
@@ -18,6 +25,21 @@
     from dreadnode.optimization.trial import Trial
 
 
+# Compliance framework tags for GOAT attack
+# Core jailbreak technique tags - specific vulnerability categories (LLM02, LLM07, etc.)
+# are added when transforms targeting those categories are used
+COMPLIANCE_TAGS = tag_attack(
+    atlas=[
+        ATLASTechnique.PROMPT_INJECTION_DIRECT,
+        ATLASTechnique.LLM_JAILBREAK,
+    ],
+    owasp=OWASPCategory.LLM01_PROMPT_INJECTION,
+    saif=SAIFCategory.INPUT_MANIPULATION,
+    nist_function=NISTAIRMFFunction.MEASURE,
+    nist_subcategory="MS-2.7",
+)
+
+
 def goat_attack(
     goal: str,
     target: "Target[DnMessage, DnMessage]",
@@ -121,6 +143,7 @@ async def message_refiner(trials: list["Trial[DnMessage]"]) -> DnMessage:
         },
         constraints=[topic_constraint],
         hooks=hooks or [],
+        compliance_tags=COMPLIANCE_TAGS,
     )
 
     if early_stopping_score is not None:

dreadnode/airt/attack/prompt.py

@@ -3,6 +3,7 @@
 import rigging as rg
 
 from dreadnode.airt.attack.base import Attack
+from dreadnode.airt.compliance import ATLASTechnique, OWASPCategory, SAIFCategory, tag_attack
 from dreadnode.data_types.message import Message as DnMessage
 from dreadnode.meta import TrialCandidate
 from dreadnode.optimization.search.graph import beam_search
@@ -18,6 +19,19 @@
     from dreadnode.optimization.trial import Trial
 
 
+# Compliance framework tags for prompt attack
+# Core jailbreak technique tags - specific vulnerability categories (LLM02, LLM07, etc.)
+# are added when transforms targeting those categories are used
+COMPLIANCE_TAGS = tag_attack(
+    atlas=[
+        ATLASTechnique.PROMPT_INJECTION_DIRECT,
+        ATLASTechnique.LLM_JAILBREAK,
+    ],
+    owasp=OWASPCategory.LLM01_PROMPT_INJECTION,
+    saif=SAIFCategory.INPUT_MANIPULATION,
+)
+
+
 def prompt_attack(
     goal: str,
     target: "Target[DnMessage, DnMessage]",
@@ -117,6 +131,7 @@ async def message_refiner(trials: list["Trial[DnMessage]"]) -> DnMessage:
             "prompt_judge": prompt_judge,
         },
         hooks=hooks or [],
+        compliance_tags=COMPLIANCE_TAGS,
     )
 
     if early_stopping_score is not None:

dreadnode/airt/attack/tap.py

@@ -2,6 +2,13 @@
 
 from dreadnode.airt.attack import Attack
 from dreadnode.airt.attack.prompt import prompt_attack
+from dreadnode.airt.compliance import (
+    ATLASTechnique,
+    NISTAIRMFFunction,
+    OWASPCategory,
+    SAIFCategory,
+    tag_attack,
+)
 from dreadnode.data_types.message import Message as DnMessage
 from dreadnode.scorers.judge import llm_judge
 
@@ -10,6 +17,21 @@
     from dreadnode.eval.hooks.base import EvalHook
 
 
+# Compliance framework tags for TAP attack
+# Core jailbreak technique tags - specific vulnerability categories (LLM02, LLM07, etc.)
+# are added when transforms targeting those categories are used
+COMPLIANCE_TAGS = tag_attack(
+    atlas=[
+        ATLASTechnique.PROMPT_INJECTION_DIRECT,
+        ATLASTechnique.LLM_JAILBREAK,
+    ],
+    owasp=OWASPCategory.LLM01_PROMPT_INJECTION,
+    saif=SAIFCategory.INPUT_MANIPULATION,
+    nist_function=NISTAIRMFFunction.MEASURE,
+    nist_subcategory="MS-2.7",
+)
+
+
 def tap_attack(
     goal: str,
     target: "Target[DnMessage, DnMessage]",
@@ -45,7 +67,7 @@ def tap_attack(
 
     topic_constraint = llm_judge(evaluator_model, ON_TOPIC_RUBRIC.format(goal=goal))
 
-    return prompt_attack(
+    base_attack = prompt_attack(
         goal,
         target,
         attacker_model,
@@ -58,7 +80,13 @@ def tap_attack(
         branching_factor=branching_factor,
         context_depth=context_depth,
         hooks=hooks or [],
-    ).with_(constraints={"on_topic": topic_constraint})
+    )
+
+    # Set compliance tags before cloning
+    base_attack.compliance_tags = COMPLIANCE_TAGS
+
+    # Add constraint and return
+    return base_attack.with_(constraints={"on_topic": topic_constraint})
 
 
 REFINE_GUIDANCE = """\