Block dangerous AI tool calls via run_vulnerability_scan.

README.md

@@ -23,6 +23,7 @@ Zen will autonomously protect your Python applications from the inside against:
 * 🛡️ [Path traversal attacks](https://www.aikido.dev/blog/path-traversal-in-2024-the-year-unpacked)
 * 🛡️ [Server-side request forgery (SSRF)](./docs/ssrf.md)
 * 🛡️ [Attack wave detection](https://help.aikido.dev/zen-firewall/zen-features/attack-wave-protection)
+* 🛡️ [Dangerous AI tool calls](./docs/ai-tool-call-blocking.md) (experimental — block host file/shell tools in AI responses)
 
 Zen operates autonomously on the same server as your Python app to:
 

aikido_zen/errors/__init__.py

@@ -61,3 +61,15 @@ class AikidoSSRF(AikidoException):
     """Exception because of SSRF"""
 
     kind = "ssrf"
+
+
+class AikidoAIToolCallBlocked(AikidoException):
+    """Exception when an AI provider response requests a blocked tool call"""
+
+    kind = "ai_tool_call"
+
+    def __init__(self, tool_name="unknown", pattern=None):
+        message = generate_default_message(self.kind) + f", tool: {tool_name}"
+        if pattern:
+            message += f", pattern: {pattern}"
+        super().__init__(message)

aikido_zen/helpers/attack_human_name.py

@@ -13,4 +13,6 @@ def attack_human_name(kind):
         return "a path traversal attack"
     if kind == "ssrf":
         return "a server-side request forgery"
+    if kind == "ai_tool_call":
+        return "a dangerous AI tool call"
     return "unknown"

aikido_zen/helpers/get_blocked_ai_tool_rules.py

@@ -0,0 +1,18 @@
+"""
+Helper function file, see function docstring
+"""
+
+import os
+
+
+def get_blocked_ai_tool_rules():
+    """
+    Reads AIKIDO_BLOCKED_AI_TOOL_NAMES and AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS.
+    Returns (tool_names: set[str], arg_patterns: list[str]).
+    """
+    names_raw = os.getenv("AIKIDO_BLOCKED_AI_TOOL_NAMES", "")
+    patterns_raw = os.getenv("AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS", "")
+
+    names = {part.strip() for part in names_raw.split(",") if part.strip()}
+    patterns = [part.strip() for part in patterns_raw.split(",") if part.strip()]
+    return names, patterns

aikido_zen/sinks/openai.py

@@ -2,8 +2,10 @@
 patching module openai
 - patches function create(...) on Responses class, to inspect response
 - patches function create(...) on Completions class, to inspect response
+- inspects tool_calls in the returned payload and can block dangerous ones
 """
 
+import aikido_zen.vulnerabilities as vulns
 from aikido_zen.helpers.on_ai_call import on_ai_call
 from aikido_zen.helpers.register_call import register_call
 from aikido_zen.sinks import on_import, patch_function, after
@@ -30,6 +32,7 @@ def _create_responses(func, instance, args, kwargs, return_value):
         input_tokens=res.usage.input_tokens,
         output_tokens=res.usage.output_tokens,
     )
+    vulns.run_vulnerability_scan(kind="ai_tool_call", op=op, args=(res,))
 
 
 @after
@@ -44,6 +47,7 @@ def _create_completions(func, instance, args, kwargs, return_value):
         input_tokens=res.usage.prompt_tokens,
         output_tokens=res.usage.completion_tokens,
     )
+    vulns.run_vulnerability_scan(kind="ai_tool_call", op=op, args=(res,))
 
 
 @on_import("openai.resources.responses.responses", "openai", "1.0")

aikido_zen/vulnerabilities/__init__.py

@@ -11,6 +11,7 @@
     AikidoShellInjection,
     AikidoPathTraversal,
     AikidoSSRF,
+    AikidoAIToolCallBlocked,
 )
 import aikido_zen.background_process.comms as comm
 from aikido_zen.helpers.logging import logger
@@ -29,6 +30,9 @@
 from .path_traversal.check_context_for_path_traversal import (
     check_context_for_path_traversal,
 )
+from .ai_tool_call.check_ai_response_for_blocked_tool_call import (
+    check_ai_response_for_blocked_tool_call,
+)
 from ..background_process.commands import PutEventCommand
 from ..helpers.create_detected_attack_api_event import create_detected_attack_api_event
 from ..helpers.ipc.send_payload import send_payload
@@ -44,14 +48,17 @@ def run_vulnerability_scan(kind, op, args):
     if is_protection_forced_off_cached(context):
         return
 
+    # SSRF and AI tool calls can be inspected without request context / thread cache.
+    kinds_without_context = ("ssrf", "ai_tool_call")
+
     comms = comm.get_comms()
     thread_cache = get_cache()
-    if not context and kind != "ssrf":
+    if not context and kind not in kinds_without_context:
         # Make a special exception for SSRF, which checks itself if context is set.
         # This is because some scans/tests for SSRF do not require a context to be set.
         return
 
-    if not thread_cache and kind != "ssrf":
+    if not thread_cache and kind not in kinds_without_context:
         # Make a special exception for SSRF, which checks itself if thread cache is set.
         # This is because some scans/tests for SSRF do not require a thread cache to be set.
         return
@@ -89,6 +96,14 @@ def run_vulnerability_scan(kind, op, args):
             dns_results, hostname, port = args
             injection_results = inspect_getaddrinfo_result(dns_results, hostname, port)
             error_type = AikidoSSRF
+        elif kind == "ai_tool_call":
+            injection_results = check_ai_response_for_blocked_tool_call(
+                response=args[0], operation=op
+            )
+            error_type = AikidoAIToolCallBlocked
+            if injection_results:
+                metadata = injection_results.get("metadata") or {}
+                error_args = (metadata.get("tool", "unknown"), metadata.get("pattern"))
         else:
             logger.error(
                 "Vulnerability type %s currently has no scans implemented", kind
@@ -99,7 +114,8 @@ def run_vulnerability_scan(kind, op, args):
     if injection_results:
         blocked = is_blocking_enabled()
         operation = injection_results["operation"]
-        thread_cache.stats.on_detected_attack(blocked, operation)
+        if thread_cache:
+            thread_cache.stats.on_detected_attack(blocked, operation)
 
         stack = get_clean_stacktrace()
 

aikido_zen/vulnerabilities/ai_tool_call/check_ai_response_for_blocked_tool_call.py

@@ -0,0 +1,43 @@
+"""
+Exports `check_ai_response_for_blocked_tool_call`
+"""
+
+from aikido_zen.helpers.get_blocked_ai_tool_rules import get_blocked_ai_tool_rules
+from .detect_blocked_ai_tool_call import (
+    argument_matches_pattern,
+    detect_blocked_ai_tool_call,
+)
+from .extract_tool_calls_from_ai_response import extract_tool_calls_from_ai_response
+
+
+def check_ai_response_for_blocked_tool_call(response, operation):
+    """
+    Inspect an AI provider response for blocked tool calls.
+    Returns an attack dict, or {} when nothing matches.
+    """
+    names, patterns = get_blocked_ai_tool_rules()
+    if not names and not patterns:
+        return {}
+
+    for tool_name, arguments in extract_tool_calls_from_ai_response(response):
+        if not detect_blocked_ai_tool_call(
+            tool_name, arguments, names=names, patterns=patterns
+        ):
+            continue
+
+        metadata = {"tool": tool_name}
+        for pattern in patterns:
+            if argument_matches_pattern(arguments, pattern):
+                metadata["pattern"] = pattern
+                break
+
+        return {
+            "operation": operation,
+            "kind": "ai_tool_call",
+            "source": "",
+            "pathToPayload": "",
+            "metadata": metadata,
+            "payload": {"tool": tool_name, "arguments": arguments},
+        }
+
+    return {}

aikido_zen/vulnerabilities/ai_tool_call/check_ai_response_for_blocked_tool_call_test.py

@@ -0,0 +1,60 @@
+"""Unit tests for check_ai_response_for_blocked_tool_call"""
+
+from types import SimpleNamespace
+
+from aikido_zen.vulnerabilities.ai_tool_call.check_ai_response_for_blocked_tool_call import (
+    check_ai_response_for_blocked_tool_call,
+)
+
+
+def _chat_completion_with_tool_call(name, arguments):
+    return SimpleNamespace(
+        choices=[
+            SimpleNamespace(
+                message=SimpleNamespace(
+                    tool_calls=[
+                        SimpleNamespace(
+                            function=SimpleNamespace(name=name, arguments=arguments)
+                        )
+                    ]
+                )
+            )
+        ]
+    )
+
+
+def test_check_returns_attack_dict_when_blocked(monkeypatch):
+    monkeypatch.setenv("AIKIDO_BLOCKED_AI_TOOL_NAMES", "read_file")
+    monkeypatch.setenv("AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS", "**/.env")
+
+    result = check_ai_response_for_blocked_tool_call(
+        _chat_completion_with_tool_call("read_file", '{"path": "config/.env"}'),
+        operation="openai.Completions.create",
+    )
+    assert result["kind"] == "ai_tool_call"
+    assert result["operation"] == "openai.Completions.create"
+    assert result["metadata"]["tool"] == "read_file"
+    assert result["metadata"]["pattern"] == "**/.env"
+    assert result["payload"]["tool"] == "read_file"
+
+
+def test_check_returns_empty_when_allowed(monkeypatch):
+    monkeypatch.setenv("AIKIDO_BLOCKED_AI_TOOL_NAMES", "read_file")
+    monkeypatch.setenv("AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS", "**/.env")
+
+    result = check_ai_response_for_blocked_tool_call(
+        _chat_completion_with_tool_call("read_file", '{"path": "README.md"}'),
+        operation="openai.Completions.create",
+    )
+    assert result == {}
+
+
+def test_check_returns_empty_without_rules(monkeypatch):
+    monkeypatch.delenv("AIKIDO_BLOCKED_AI_TOOL_NAMES", raising=False)
+    monkeypatch.delenv("AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS", raising=False)
+
+    result = check_ai_response_for_blocked_tool_call(
+        _chat_completion_with_tool_call("read_file", '{"path": "config/.env"}'),
+        operation="openai.Completions.create",
+    )
+    assert result == {}

aikido_zen/vulnerabilities/ai_tool_call/detect_blocked_ai_tool_call.py

@@ -0,0 +1,72 @@
+"""
+Exports `detect_blocked_ai_tool_call`
+"""
+
+import json
+from pathlib import PurePosixPath
+
+from aikido_zen.helpers.get_blocked_ai_tool_rules import get_blocked_ai_tool_rules
+
+
+def _argument_values(arguments):
+    if arguments is None:
+        return []
+    if isinstance(arguments, str):
+        values = [arguments]
+        try:
+            arguments = json.loads(arguments)
+        except (json.JSONDecodeError, TypeError):
+            return values
+    else:
+        values = []
+
+    def walk(node):
+        if isinstance(node, dict):
+            for value in node.values():
+                walk(value)
+        elif isinstance(node, list):
+            for value in node:
+                walk(value)
+        elif node is not None:
+            values.append(str(node))
+
+    if isinstance(arguments, (dict, list)):
+        walk(arguments)
+    elif not values:
+        values.append(str(arguments))
+    return values
+
+
+def argument_matches_pattern(arguments, pattern):
+    """True if any tool argument value matches the glob pattern."""
+    for value in _argument_values(arguments):
+        normalized = value.replace("\\", "/").lstrip("/")
+        try:
+            if PurePosixPath(normalized).match(pattern):
+                return True
+        except ValueError:
+            continue
+    return False
+
+
+def detect_blocked_ai_tool_call(tool_name, arguments, names=None, patterns=None):
+    """
+    Detects if a tool call matches the configured block rules.
+    """
+    if names is None or patterns is None:
+        env_names, env_patterns = get_blocked_ai_tool_rules()
+        if names is None:
+            names = env_names
+        if patterns is None:
+            patterns = env_patterns
+
+    if not names and not patterns:
+        return False
+    if names and tool_name not in names:
+        return False
+    if not patterns:
+        return True
+
+    return any(
+        argument_matches_pattern(arguments, pattern) for pattern in patterns
+    )