From c7e822ac18cbead631c40be16c43b8329f70320f Mon Sep 17 00:00:00 2001 From: Raja Sekhar Rao Dheekonda Date: Tue, 15 Sep 2026 15:33:34 -0700 Subject: [PATCH] feat(ai-red-teaming): extensible guardrail policy table (provision + teardown) Refactor the provisioning guardrail into a small policy table so new deny/redirect rules are one line. Add a teardown rule: shell `dreadnode|dn env teardown/delete/ destroy` is denied and redirected to the teardown_environment tool (destructive, tool-only per the skills). Provision subcommands redirect to provision_environment; a generic `dreadnode env` catch-all redirects to the environment tools. Scope stays narrow - legitimate CLI like `dn airt run` is unaffected (no skill forbids it). Bump to 1.16.3. --- capabilities/ai-red-teaming/capability.yaml | 2 +- .../hooks/block_cli_provisioning.py | 89 ++++++++++++++----- .../tests/test_block_cli_provisioning.py | 13 +++ 3 files changed, 79 insertions(+), 25 deletions(-) diff --git a/capabilities/ai-red-teaming/capability.yaml b/capabilities/ai-red-teaming/capability.yaml index 8bf39a3..a1e1713 100644 --- a/capabilities/ai-red-teaming/capability.yaml +++ b/capabilities/ai-red-teaming/capability.yaml @@ -1,6 +1,6 @@ schema: 1 name: ai-red-teaming -version: "1.16.2" +version: "1.16.3" description: > Probe the security and safety of AI applications, agents, and foundation models. Orchestrates adversarial attack workflows to discover vulnerabilities in LLMs, diff --git a/capabilities/ai-red-teaming/hooks/block_cli_provisioning.py b/capabilities/ai-red-teaming/hooks/block_cli_provisioning.py index 1abb5a7..bcd2022 100644 --- a/capabilities/ai-red-teaming/hooks/block_cli_provisioning.py +++ b/capabilities/ai-red-teaming/hooks/block_cli_provisioning.py @@ -1,12 +1,19 @@ -"""Structural guardrail: deny environment provisioning through the shell. - -The `provisioning-and-lifecycle` skill instructs the agent to provision only via -the `provision_environment` tool and to NEVER run `dreadnode env ...` in a shell. -Skill prose is advisory - a weaker driver model can ignore it and fall back to -the CLI (ENG-8477). This hook enforces the rule structurally: if the agent tries -to provision an environment through a shell tool, the call is denied before it -runs and the agent is told to use `provision_environment` (or stop and inform -the user). +"""Structural guardrails: deny shell fallbacks that bypass first-class tools. + +The AIRT skills mandate tool-only paths for a few operations - most importantly, +`provisioning-and-lifecycle` and `error-troubleshooting` both say to NEVER run +`dreadnode env ...` in a shell and to use the `provision_environment` / +`teardown_environment` tools instead. Skill prose is advisory; a weaker driver +model can ignore it and fall back to the CLI (ENG-8477). + +This module enforces those rules structurally with a small policy table. A +`ToolStart` hook inspects every shell tool call and, when its command matches a +rule, denies the call before it runs and redirects the agent to the correct +tool (via `RetryWithFeedback`, which the runtime turns into a `[POLICY DENIED]` +tool result). Rules are scoped narrowly so legitimate shell and CLI usage +(for example `dn airt run`, which the skills permit) is never affected. + +To add a guardrail, append a `_Rule` - one line per policy. """ from __future__ import annotations @@ -14,23 +21,49 @@ import json import re import typing as t +from dataclasses import dataclass from dreadnode.agents.events import ToolStart from dreadnode.agents.reactions import RetryWithFeedback from dreadnode.core.hook import hook # Shell-style tools that can execute a CLI command. -_SHELL_TOOLS = {"bash", "shell", "sh", "python", "python3", "run_command", "execute"} +_SHELL_TOOLS = frozenset({"bash", "shell", "sh", "python", "python3", "run_command", "execute"}) + + +@dataclass(frozen=True) +class _Rule: + """A single deny-and-redirect policy for shell tool calls.""" -# `dreadnode env ...` / `dn env ...` / `dreadnode environment ...` in a shell. -# Matches the exact pattern the skill forbids without touching other CLI usage. -_CLI_PROVISION = re.compile(r"\b(?:dreadnode|dn)\s+env(?:ironment)?\b", re.IGNORECASE) + tools: frozenset[str] + pattern: re.Pattern[str] + feedback: str -_FEEDBACK = ( - "Provisioning an environment through the shell/CLI is not permitted. " - "Use the `provision_environment` tool instead - it resolves the catalog and " - "credentials for you. If `provision_environment` is unavailable, STOP and " - "inform the user; do not fall back to the shell." + +# NOTE: order matters - the first matching rule wins, so list specific rules +# (provision / teardown) before the generic `dreadnode env` catch-all. +_RULES: tuple[_Rule, ...] = ( + _Rule( + _SHELL_TOOLS, + re.compile(r"\b(?:dreadnode|dn)\s+env(?:ironment)?\s+(?:provision|create|up|start|new)\b", re.IGNORECASE), + "Provisioning an environment through the shell/CLI is not permitted. Use the " + "`provision_environment` tool instead - it resolves the catalog and credentials for you. " + "If it is unavailable, STOP and inform the user; do not fall back to the shell.", + ), + _Rule( + _SHELL_TOOLS, + re.compile(r"\b(?:dreadnode|dn)\s+env(?:ironment)?\s+(?:teardown|delete|destroy|down|stop|rm|remove)\b", re.IGNORECASE), + "Tearing down an environment through the shell/CLI is not permitted. Use the " + "`teardown_environment` tool instead. If it is unavailable, STOP and inform the user; " + "do not fall back to the shell.", + ), + _Rule( + _SHELL_TOOLS, + re.compile(r"\b(?:dreadnode|dn)\s+env(?:ironment)?\b", re.IGNORECASE), + "Managing environments through the shell/CLI is not permitted. Use the environment tools " + "(`provision_environment`, `list_environments`, `teardown_environment`) instead. If they " + "are unavailable, STOP and inform the user; do not fall back to the shell.", + ), ) @@ -46,19 +79,27 @@ def _iter_strings(value: t.Any) -> t.Iterator[str]: yield from _iter_strings(item) +def _match(tool_name: str, args: t.Any) -> _Rule | None: + """Return the first rule that denies this tool call, if any.""" + name = tool_name.lower() + texts = list(_iter_strings(args)) + for rule in _RULES: + if name in rule.tools and any(rule.pattern.search(text) for text in texts): + return rule + return None + + @hook(ToolStart) async def block_cli_provisioning(event: ToolStart) -> RetryWithFeedback | None: - """Deny shell-based environment provisioning and redirect to the tool.""" - if event.tool_call.name.lower() not in _SHELL_TOOLS: - return None - + """Deny shell calls that bypass first-class tools and redirect the agent.""" raw = event.tool_call.function.arguments try: args: t.Any = json.loads(raw) except (json.JSONDecodeError, TypeError): args = raw - if any(_CLI_PROVISION.search(text) for text in _iter_strings(args)): - return RetryWithFeedback(feedback=_FEEDBACK, tool_call_id=event.tool_call.id) + rule = _match(event.tool_call.name, args) + if rule is not None: + return RetryWithFeedback(feedback=rule.feedback, tool_call_id=event.tool_call.id) return None diff --git a/capabilities/ai-red-teaming/tests/test_block_cli_provisioning.py b/capabilities/ai-red-teaming/tests/test_block_cli_provisioning.py index d4efa25..bf35fc8 100644 --- a/capabilities/ai-red-teaming/tests/test_block_cli_provisioning.py +++ b/capabilities/ai-red-teaming/tests/test_block_cli_provisioning.py @@ -5,12 +5,16 @@ import asyncio import importlib.util import json +import sys from pathlib import Path _HOOK_PATH = Path(__file__).resolve().parents[1] / "hooks" / "block_cli_provisioning.py" _spec = importlib.util.spec_from_file_location("airt_block_cli_provisioning", _HOOK_PATH) assert _spec and _spec.loader _mod = importlib.util.module_from_spec(_spec) +# Register before exec so dataclasses can resolve the module namespace under +# `from __future__ import annotations` (dataclasses looks the module up in sys.modules). +sys.modules[_spec.name] = _mod _spec.loader.exec_module(_mod) @@ -51,8 +55,17 @@ def test_blocks_dn_env_and_environment_variants() -> None: assert _run("python", json.dumps({"code": "os.system('dreadnode env provision m')"})) is not None +def test_blocks_teardown_via_cli_and_redirects_to_teardown_tool() -> None: + r = _run("bash", json.dumps({"command": "dreadnode env teardown finops-mesh"})) + assert r is not None + assert "teardown_environment" in r.feedback + # destroy/delete variants also redirect to the teardown tool + assert "teardown_environment" in _run("bash", json.dumps({"command": "dn env delete soc-mesh"})).feedback + + def test_allows_non_provisioning_shell_commands() -> None: assert _run("bash", json.dumps({"command": "ls -la /home/user"})) is None + # dn airt run is a permitted CLI (no skill forbids it) - must not be blocked assert _run("bash", json.dumps({"command": "dn airt run --goal x --attack tap"})) is None assert _run("bash", json.dumps({"command": "python attack.py"})) is None