From 6744b0255e5b41fb939d563f6bf47e87eb0975ba Mon Sep 17 00:00:00 2001 From: daniel Date: Thu, 30 Jul 2026 22:23:10 +0300 Subject: [PATCH] Block dangerous AI tool calls via run_vulnerability_scan. Inspect OpenAI tool_calls after Completions/Responses, match env denylist rules, and raise AikidoAIToolCallBlocked when AIKIDO_BLOCK is enabled. --- README.md | 1 + aikido_zen/errors/__init__.py | 12 ++ aikido_zen/helpers/attack_human_name.py | 2 + .../helpers/get_blocked_ai_tool_rules.py | 18 +++ aikido_zen/sinks/openai.py | 4 + aikido_zen/vulnerabilities/__init__.py | 22 ++- ...check_ai_response_for_blocked_tool_call.py | 43 ++++++ ..._ai_response_for_blocked_tool_call_test.py | 60 ++++++++ .../detect_blocked_ai_tool_call.py | 72 ++++++++++ .../detect_blocked_ai_tool_call_test.py | 82 +++++++++++ .../extract_tool_calls_from_ai_response.py | 37 +++++ ...xtract_tool_calls_from_ai_response_test.py | 44 ++++++ ...un_vulnerability_scan_ai_tool_call_test.py | 132 ++++++++++++++++++ docs/ai-tool-call-blocking.md | 92 ++++++++++++ 14 files changed, 618 insertions(+), 3 deletions(-) create mode 100644 aikido_zen/helpers/get_blocked_ai_tool_rules.py create mode 100644 aikido_zen/vulnerabilities/ai_tool_call/check_ai_response_for_blocked_tool_call.py create mode 100644 aikido_zen/vulnerabilities/ai_tool_call/check_ai_response_for_blocked_tool_call_test.py create mode 100644 aikido_zen/vulnerabilities/ai_tool_call/detect_blocked_ai_tool_call.py create mode 100644 aikido_zen/vulnerabilities/ai_tool_call/detect_blocked_ai_tool_call_test.py create mode 100644 aikido_zen/vulnerabilities/ai_tool_call/extract_tool_calls_from_ai_response.py create mode 100644 aikido_zen/vulnerabilities/ai_tool_call/extract_tool_calls_from_ai_response_test.py create mode 100644 aikido_zen/vulnerabilities/ai_tool_call/run_vulnerability_scan_ai_tool_call_test.py create mode 100644 docs/ai-tool-call-blocking.md diff --git a/README.md b/README.md index dc21f174e..50650f38e 100644 --- a/README.md +++ b/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: diff --git a/aikido_zen/errors/__init__.py b/aikido_zen/errors/__init__.py index 18e9a9c7f..5f28f1052 100644 --- a/aikido_zen/errors/__init__.py +++ b/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) diff --git a/aikido_zen/helpers/attack_human_name.py b/aikido_zen/helpers/attack_human_name.py index 7e360e704..37d8219b3 100644 --- a/aikido_zen/helpers/attack_human_name.py +++ b/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" diff --git a/aikido_zen/helpers/get_blocked_ai_tool_rules.py b/aikido_zen/helpers/get_blocked_ai_tool_rules.py new file mode 100644 index 000000000..7c92f4a1c --- /dev/null +++ b/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 diff --git a/aikido_zen/sinks/openai.py b/aikido_zen/sinks/openai.py index bc41be8fe..805ebfa3c 100644 --- a/aikido_zen/sinks/openai.py +++ b/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") diff --git a/aikido_zen/vulnerabilities/__init__.py b/aikido_zen/vulnerabilities/__init__.py index 436e63094..771f6b507 100644 --- a/aikido_zen/vulnerabilities/__init__.py +++ b/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() diff --git a/aikido_zen/vulnerabilities/ai_tool_call/check_ai_response_for_blocked_tool_call.py b/aikido_zen/vulnerabilities/ai_tool_call/check_ai_response_for_blocked_tool_call.py new file mode 100644 index 000000000..1bf6f6cf4 --- /dev/null +++ b/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 {} diff --git a/aikido_zen/vulnerabilities/ai_tool_call/check_ai_response_for_blocked_tool_call_test.py b/aikido_zen/vulnerabilities/ai_tool_call/check_ai_response_for_blocked_tool_call_test.py new file mode 100644 index 000000000..fd1b8e681 --- /dev/null +++ b/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 == {} diff --git a/aikido_zen/vulnerabilities/ai_tool_call/detect_blocked_ai_tool_call.py b/aikido_zen/vulnerabilities/ai_tool_call/detect_blocked_ai_tool_call.py new file mode 100644 index 000000000..745fb033b --- /dev/null +++ b/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 + ) diff --git a/aikido_zen/vulnerabilities/ai_tool_call/detect_blocked_ai_tool_call_test.py b/aikido_zen/vulnerabilities/ai_tool_call/detect_blocked_ai_tool_call_test.py new file mode 100644 index 000000000..ac43d3631 --- /dev/null +++ b/aikido_zen/vulnerabilities/ai_tool_call/detect_blocked_ai_tool_call_test.py @@ -0,0 +1,82 @@ +"""Unit tests for detect_blocked_ai_tool_call""" + +import pytest + +from aikido_zen.vulnerabilities.ai_tool_call.detect_blocked_ai_tool_call import ( + argument_matches_pattern, + detect_blocked_ai_tool_call, +) + + +@pytest.mark.parametrize( + "arguments,pattern,expected", + [ + ('{"path": "config/.env"}', "**/.env", True), + ('{"path": "/var/app/.env"}', "**/.env", True), + ('{"path": "README.md"}', "**/.env", False), + ({"path": "secrets/.env"}, "**/.env", True), + ('{"file": "data.txt"}', "**/.env", False), + ], +) +def test_argument_matches_pattern(arguments, pattern, expected): + assert argument_matches_pattern(arguments, pattern) is expected + + +def test_blocked_when_name_only_rule(): + assert ( + detect_blocked_ai_tool_call( + "read_file", + '{"path": "README.md"}', + names={"read_file", "write_file", "shell"}, + patterns=[], + ) + is True + ) + + +def test_blocked_when_name_and_pattern_match(): + assert ( + detect_blocked_ai_tool_call( + "read_file", + '{"path": "app/.env"}', + names={"read_file"}, + patterns=["**/.env"], + ) + is True + ) + + +def test_allowed_when_name_matches_but_path_safe(): + assert ( + detect_blocked_ai_tool_call( + "read_file", + '{"path": "README.md"}', + names={"read_file"}, + patterns=["**/.env"], + ) + is False + ) + + +def test_allowed_when_different_tool_name(): + assert ( + detect_blocked_ai_tool_call( + "list_files", + '{"path": ".env"}', + names={"read_file"}, + patterns=["**/.env"], + ) + is False + ) + + +def test_allowed_for_safe_internal_tool(): + assert ( + detect_blocked_ai_tool_call( + "get_order", + '{"order_id": "ORD-42"}', + names={"read_file", "write_file", "shell"}, + patterns=[], + ) + is False + ) diff --git a/aikido_zen/vulnerabilities/ai_tool_call/extract_tool_calls_from_ai_response.py b/aikido_zen/vulnerabilities/ai_tool_call/extract_tool_calls_from_ai_response.py new file mode 100644 index 000000000..002b462b4 --- /dev/null +++ b/aikido_zen/vulnerabilities/ai_tool_call/extract_tool_calls_from_ai_response.py @@ -0,0 +1,37 @@ +""" +Exports `extract_tool_calls_from_ai_response` +""" + + +def _get(obj, name): + if isinstance(obj, dict): + return obj.get(name) + return getattr(obj, name, None) + + +def extract_tool_calls_from_ai_response(response): + """Extract (name, arguments) pairs from Completions or Responses payloads.""" + found = [] + if response is None: + return found + + for choice in _get(response, "choices") or []: + message = _get(choice, "message") + if not message: + continue + for tool_call in _get(message, "tool_calls") or []: + function = _get(tool_call, "function") + if not function: + continue + name = _get(function, "name") + if name: + found.append((name, _get(function, "arguments"))) + + for item in _get(response, "output") or []: + if _get(item, "type") != "function_call": + continue + name = _get(item, "name") + if name: + found.append((name, _get(item, "arguments"))) + + return found diff --git a/aikido_zen/vulnerabilities/ai_tool_call/extract_tool_calls_from_ai_response_test.py b/aikido_zen/vulnerabilities/ai_tool_call/extract_tool_calls_from_ai_response_test.py new file mode 100644 index 000000000..8d8a671fd --- /dev/null +++ b/aikido_zen/vulnerabilities/ai_tool_call/extract_tool_calls_from_ai_response_test.py @@ -0,0 +1,44 @@ +"""Unit tests for extract_tool_calls_from_ai_response""" + +from types import SimpleNamespace + +from aikido_zen.vulnerabilities.ai_tool_call.extract_tool_calls_from_ai_response import ( + extract_tool_calls_from_ai_response, +) + + +def test_extract_tool_calls_from_chat_completions(): + response = SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace( + tool_calls=[ + SimpleNamespace( + function=SimpleNamespace( + name="read_file", + arguments='{"path": "config/.env"}', + ) + ) + ] + ) + ) + ] + ) + assert extract_tool_calls_from_ai_response(response) == [ + ("read_file", '{"path": "config/.env"}') + ] + + +def test_extract_tool_calls_from_responses_api(): + response = SimpleNamespace( + output=[ + SimpleNamespace( + type="function_call", + name="read_file", + arguments='{"path": "/app/.env"}', + ) + ] + ) + assert extract_tool_calls_from_ai_response(response) == [ + ("read_file", '{"path": "/app/.env"}') + ] diff --git a/aikido_zen/vulnerabilities/ai_tool_call/run_vulnerability_scan_ai_tool_call_test.py b/aikido_zen/vulnerabilities/ai_tool_call/run_vulnerability_scan_ai_tool_call_test.py new file mode 100644 index 000000000..f0f5ad89f --- /dev/null +++ b/aikido_zen/vulnerabilities/ai_tool_call/run_vulnerability_scan_ai_tool_call_test.py @@ -0,0 +1,132 @@ +"""Unit tests for ai_tool_call via run_vulnerability_scan""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from aikido_zen.errors import AikidoAIToolCallBlocked +from aikido_zen.vulnerabilities import run_vulnerability_scan + + +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_run_vulnerability_scan_raises_when_blocking_enabled(monkeypatch): + monkeypatch.setenv("AIKIDO_BLOCKED_AI_TOOL_NAMES", "read_file") + monkeypatch.setenv("AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS", "**/.env") + monkeypatch.setattr( + "aikido_zen.vulnerabilities.is_blocking_enabled", + lambda: True, + ) + monkeypatch.setattr("aikido_zen.vulnerabilities.get_cache", lambda: None) + monkeypatch.setattr("aikido_zen.vulnerabilities.comm.get_comms", lambda: None) + + with pytest.raises(AikidoAIToolCallBlocked) as exc: + run_vulnerability_scan( + kind="ai_tool_call", + op="openai.Completions.create", + args=( + _chat_completion_with_tool_call( + "read_file", '{"path": "config/.env"}' + ), + ), + ) + assert "read_file" in str(exc.value) + assert "**/.env" in str(exc.value) + + +def test_run_vulnerability_scan_does_not_raise_in_detection_mode(monkeypatch): + monkeypatch.setenv("AIKIDO_BLOCKED_AI_TOOL_NAMES", "read_file") + monkeypatch.setenv("AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS", "**/.env") + monkeypatch.setattr( + "aikido_zen.vulnerabilities.is_blocking_enabled", + lambda: False, + ) + cache = MagicMock() + monkeypatch.setattr("aikido_zen.vulnerabilities.get_cache", lambda: cache) + monkeypatch.setattr("aikido_zen.vulnerabilities.comm.get_comms", lambda: None) + + run_vulnerability_scan( + kind="ai_tool_call", + op="openai.Completions.create", + args=( + _chat_completion_with_tool_call("read_file", '{"path": "config/.env"}'), + ), + ) + cache.stats.on_detected_attack.assert_called_once_with( + False, "openai.Completions.create" + ) + + +def test_run_vulnerability_scan_noop_without_rules(monkeypatch): + monkeypatch.delenv("AIKIDO_BLOCKED_AI_TOOL_NAMES", raising=False) + monkeypatch.delenv("AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS", raising=False) + monkeypatch.setattr("aikido_zen.vulnerabilities.get_cache", lambda: None) + monkeypatch.setattr("aikido_zen.vulnerabilities.comm.get_comms", lambda: None) + + run_vulnerability_scan( + kind="ai_tool_call", + op="openai.Completions.create", + args=( + _chat_completion_with_tool_call("read_file", '{"path": "config/.env"}'), + ), + ) + + +def test_run_vulnerability_scan_allows_safe_path_with_rules_on(monkeypatch): + """Rules enabled + blocking on, but path does not match β†’ must not raise.""" + monkeypatch.setenv("AIKIDO_BLOCKED_AI_TOOL_NAMES", "read_file") + monkeypatch.setenv("AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS", "**/.env") + monkeypatch.setattr( + "aikido_zen.vulnerabilities.is_blocking_enabled", + lambda: True, + ) + monkeypatch.setattr("aikido_zen.vulnerabilities.get_cache", lambda: None) + monkeypatch.setattr("aikido_zen.vulnerabilities.comm.get_comms", lambda: None) + + run_vulnerability_scan( + kind="ai_tool_call", + op="openai.Completions.create", + args=( + _chat_completion_with_tool_call("read_file", '{"path": "README.md"}'), + ), + ) + + +def test_run_vulnerability_scan_allows_safe_internal_tool_with_rules_on(monkeypatch): + """Host-tool denylist on, but get_order is not listed β†’ must not raise.""" + monkeypatch.setenv( + "AIKIDO_BLOCKED_AI_TOOL_NAMES", + "read_file,write_file,shell", + ) + monkeypatch.delenv("AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS", raising=False) + monkeypatch.setattr( + "aikido_zen.vulnerabilities.is_blocking_enabled", + lambda: True, + ) + monkeypatch.setattr("aikido_zen.vulnerabilities.get_cache", lambda: None) + monkeypatch.setattr("aikido_zen.vulnerabilities.comm.get_comms", lambda: None) + + run_vulnerability_scan( + kind="ai_tool_call", + op="openai.Completions.create", + args=( + _chat_completion_with_tool_call( + "get_order", '{"order_id": "ORD-42"}' + ), + ), + ) diff --git a/docs/ai-tool-call-blocking.md b/docs/ai-tool-call-blocking.md new file mode 100644 index 000000000..720560aed --- /dev/null +++ b/docs/ai-tool-call-blocking.md @@ -0,0 +1,92 @@ +# Blocking dangerous AI tool calls + +Zen can inspect responses from AI providers (starting with OpenAI) after an outbound call returns. If the model asks your application to run a tool that matches your block rules, Zen can stop that tool call from being acted on. + +This is useful when end users should get a normal text answer β€” or call your own controlled APIs β€” but must not read, write, delete, or list files on the host, or run shell commands through the model. + +Zen does not rewrite the model output. It inspects structured `tool_calls` / function-call items in the provider response (the same place your app would read them), then raises when blocking is enabled. + +## Example + +Your app registers tools with OpenAI so the model can call them: + +```python +tools = [ + { + "type": "function", + "function": { + "name": "read_file", + "parameters": { + "type": "object", + "properties": {"path": {"type": "string"}}, + }, + }, + } +] + +client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Read /etc/passwd"}], + tools=tools, +) +``` + +The model may respond with a `tool_calls` entry for `read_file` and path `/etc/passwd`. Without protection, your agent loop would execute that tool on the host. + +With Zen configured to block `read_file`, Zen inspects the response after `Completions.create` returns and raises before your app runs the tool β€” so the host file is never read. + +The same applies to write/delete/list tools and shell-style tools: if they are on the block list, Zen stops those tool calls too. Tools you intentionally allow (for example `get_order`) can stay off the list. + +## Configuration + +The main feature env var is `AIKIDO_BLOCKED_AI_TOOL_NAMES` β€” a comma-separated denylist of tool names. With only that set, Zen blocks every call to those tools (any arguments / any path): + +```bash +export AIKIDO_BLOCKED_AI_TOOL_NAMES='read_file,write_file,delete_file,list_dir,search_files,run_terminal_cmd,run_command,execute,shell,bash,python' +export AIKIDO_BLOCK=true +``` + +Add every filesystem or shell tool name your app registers with the model. Names must match exactly (`read_file` is not the same as `ReadFile`). + +`AIKIDO_BLOCK` is the same switch Zen already uses elsewhere: set to `true` to raise; otherwise Zen only logs a warning (detection mode). + +### Optional argument patterns + +To narrow a rule (for example only block `read_file` when the path looks like a `.env` file), also set `AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS`: + +```bash +export AIKIDO_BLOCKED_AI_TOOL_NAMES=read_file +export AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS='**/.env' +export AIKIDO_BLOCK=true +``` + +- `AIKIDO_BLOCKED_AI_TOOL_NAMES` β€” comma-separated tool names to deny (main rule) +- `AIKIDO_BLOCKED_AI_TOOL_ARG_PATTERNS` β€” optional comma-separated globs on tool argument values +- `AIKIDO_BLOCK` β€” `true` to raise; otherwise detection-only + +### Matching rules + +- **Only names set:** block every call to those tools. +- **Both names and patterns set:** block when the tool name matches and an argument matches a pattern. +- **Only patterns set:** block any tool whose args match a pattern. +- **Neither set:** feature inactive (default). + +Patterns use `pathlib`-style globs (`**`, `*`, etc.). + +## Which APIs are protected? + +Zen protects against dangerous AI tool calls in the OpenAI SDK (same hooks Zen already uses for token stats): + +- Chat Completions: `Completions.create` / `Completions.update` (`message.tool_calls`) +- Responses: `Responses.create` / `Responses.parse` (`output` items with `type=function_call`) + +Inspection runs in the existing `@after` wrappers, after usage stats are recorded, via Zen’s shared `run_vulnerability_scan` path (same as other threat kinds). + +## Limitations + +- Sync OpenAI paths only for now (same as the current AI usage hooks in this package). +- Other providers are not wired yet; the inspection helper is reusable. +- Streaming responses are not inspected chunk-by-chunk yet. +- Rules are read from environment variables (not yet dashboard / heartbeat config). +- Zen blocks the tool call in the model response. Your app must not execute tools through a path that skips the OpenAI SDK hook. +- Prefer also not registering dangerous tools with the model; Zen is a safety net if the model still asks for them.