diff --git a/cookbook/atr_detection_callback/README.md b/cookbook/atr_detection_callback/README.md new file mode 100644 index 00000000000..fe6030b1356 --- /dev/null +++ b/cookbook/atr_detection_callback/README.md @@ -0,0 +1,44 @@ +# Agent Threat Rules detection callback + +A small cookbook example that screens user input against ATR-inspired +threat patterns before the request reaches the model. ATR is an open +detection standard for AI agent threats (prompt injection, tool +poisoning, MCP attacks, skill compromise) published under Apache-2.0: + +https://github.com/Agent-Threat-Rule/agent-threat-rules + +## What it does + +`atr_detection_callback.py` defines `ATRDetectionGuardrail`, a +`CustomGuardrail` whose `async_pre_call_hook` runs each user message +through a handful of compiled regex patterns. On a match, it logs the +rule id and raises a `ValueError`, which LiteLLM surfaces to the caller +as a blocked request. + +The patterns embedded in the file are illustrative copies covering +common categories: instruction override, system prompt exfiltration, +role-play jailbreak, base64-wrapped payloads, MCP tool override, and +`file://` SSRF references. The full ruleset lives in the ATR repository. + +## Wire it up + +Add a guardrail entry to your proxy config: + +```yaml +guardrails: + - guardrail_name: "atr-input-screen" + litellm_params: + guardrail: cookbook.atr_detection_callback.atr_detection_callback.ATRDetectionGuardrail + mode: "pre_call" + default_on: true +``` + +Then start the proxy as usual. Requests containing matching patterns +will be rejected before the LLM call, and the rule id will appear in +the proxy logs. + +## Extending + +To run against the live ATR YAML ruleset instead of embedded patterns, +load rule files at startup and compile their `detection.regex_patterns` +fields into the same list shape. diff --git a/cookbook/atr_detection_callback/atr_detection_callback.py b/cookbook/atr_detection_callback/atr_detection_callback.py new file mode 100644 index 00000000000..d4c116c62f0 --- /dev/null +++ b/cookbook/atr_detection_callback/atr_detection_callback.py @@ -0,0 +1,113 @@ +""" +ATR detection callback for LiteLLM. + +A minimal CustomGuardrail that screens user input against a small set of +ATR-inspired regex patterns covering common AI agent threats: prompt +injection overrides, role-play jailbreaks, base64-wrapped instructions, +MCP tool override, and file:// SSRF. The patterns below are illustrative +copies; the full open detection set lives in Agent Threat Rules: +https://github.com/Agent-Threat-Rule/agent-threat-rules (Apache-2.0). + +Wire it up via proxy_config.yaml: + + guardrails: + - guardrail_name: "atr-input-screen" + litellm_params: + guardrail: cookbook.atr_detection_callback.atr_detection_callback.ATRDetectionGuardrail + mode: "pre_call" + default_on: true +""" + +import re +from typing import Optional, Union + +from litellm._logging import verbose_proxy_logger +from litellm.caching.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral + +# (rule_id, label, compiled_pattern). rule_id mirrors the ATR namespace. +ATR_INSPIRED_PATTERNS = [ + ( + "ATR-PI-001", + "instruction override", + re.compile( + r"\b(ignore|disregard|forget)\s+(all\s+)?(previous|prior|above)\s+" + r"(instructions?|prompts?|rules?)", + re.IGNORECASE, + ), + ), + ( + "ATR-PI-002", + "system prompt exfiltration", + re.compile( + r"(reveal|print|repeat|show)\s+(your\s+)?" + r"(system\s+prompt|initial\s+instructions)", + re.IGNORECASE, + ), + ), + ( + "ATR-PI-003", + "role-play jailbreak", + re.compile( + r"\b(you\s+are\s+now|act\s+as|pretend\s+to\s+be)\s+" + r"(DAN|developer\s+mode|jailbroken|an?\s+unrestricted)", + re.IGNORECASE, + ), + ), + ( + "ATR-PI-004", + "base64-wrapped payload hint", + re.compile( + r"(decode|run|execute)\s+(this\s+)?base64[:\s]+[A-Za-z0-9+/=]{40,}", + re.IGNORECASE, + ), + ), + ( + "ATR-MCP-001", + "mcp tool override", + re.compile(r"<\s*(tool_override|mcp_override|new_tool_definition)\s*>", re.IGNORECASE), + ), + ( + "ATR-SSRF-001", + "file:// scheme reference", + re.compile(r"file://[^\s\"'<>]+", re.IGNORECASE), + ), +] + + +class ATRDetectionGuardrail(CustomGuardrail): + """Block requests that hit any ATR-inspired threat pattern.""" + + def __init__(self, **kwargs): + self.optional_params = kwargs + super().__init__(**kwargs) + + @staticmethod + def _scan(text: str): + for rule_id, label, pattern in ATR_INSPIRED_PATTERNS: + if pattern.search(text): + return rule_id, label + return None + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> Optional[Union[Exception, str, dict]]: + for message in data.get("messages") or []: + content = message.get("content") + if not isinstance(content, str): + continue + hit = self._scan(content) + if hit is None: + continue + rule_id, label = hit + verbose_proxy_logger.warning( + "ATR threat pattern matched: rule_id=%s label=%s", rule_id, label + ) + raise ValueError(f"Request blocked by ATR guardrail: {rule_id} ({label}).") + return data