|
| 1 | +"""AxonFlow ComputerUseGovernor — governance middleware for Anthropic Computer Use. |
| 2 | +
|
| 3 | +Evaluates Computer Use tool_use blocks against AxonFlow policies before |
| 4 | +execution, and scans tool results before feeding them back to Claude. |
| 5 | +
|
| 6 | +Drop into the sampling loop at two points: |
| 7 | +
|
| 8 | +1. After Claude returns tool_use blocks, call ``check_tool_use()`` on each |
| 9 | + block before executing it. |
| 10 | +2. After executing a tool, call ``check_result()`` on the result before |
| 11 | + appending it to the message history. |
| 12 | +
|
| 13 | +Example:: |
| 14 | +
|
| 15 | + from anthropic import Anthropic |
| 16 | + from axonflow import AxonFlow |
| 17 | + from axonflow.adapters import ComputerUseGovernor |
| 18 | +
|
| 19 | + async with AxonFlow(endpoint="http://localhost:8080") as client: |
| 20 | + governor = ComputerUseGovernor(client) |
| 21 | +
|
| 22 | + # In your sampling loop: |
| 23 | + for block in response.content: |
| 24 | + if block.type == "tool_use": |
| 25 | + check = await governor.check_tool_use({ |
| 26 | + "type": "tool_use", |
| 27 | + "id": block.id, |
| 28 | + "name": block.name, |
| 29 | + "input": block.input, |
| 30 | + }) |
| 31 | + if not check.allowed: |
| 32 | + # Skip this tool call, tell Claude it was blocked |
| 33 | + continue |
| 34 | +
|
| 35 | + result = execute_tool(block) |
| 36 | +
|
| 37 | + result_check = await governor.check_result(block.name, result) |
| 38 | + if result_check.redacted_result is not None: |
| 39 | + result = result_check.redacted_result |
| 40 | +""" |
| 41 | + |
| 42 | +from __future__ import annotations |
| 43 | + |
| 44 | +import json |
| 45 | +import re |
| 46 | +from dataclasses import dataclass |
| 47 | +from typing import TYPE_CHECKING, Any |
| 48 | + |
| 49 | +if TYPE_CHECKING: |
| 50 | + from axonflow import AxonFlow |
| 51 | + |
| 52 | + |
| 53 | +# Default patterns for dangerous bash commands |
| 54 | +DEFAULT_BLOCKED_BASH_PATTERNS: list[str] = [ |
| 55 | + r"rm\s+-rf\s+/", |
| 56 | + r"dd\s+if=", |
| 57 | + r"mkfs\b", |
| 58 | + r"curl\s+.*\|\s*(ba)?sh", |
| 59 | + r"wget\s+.*\|\s*(ba)?sh", |
| 60 | + r"cat\s+~/\.ssh/", |
| 61 | + r"cat\s+~/\.aws/", |
| 62 | + r"cat\s+\.env\b", |
| 63 | + r"chmod\s+777\b", |
| 64 | + r">\s*/dev/sd[a-z]", |
| 65 | +] |
| 66 | + |
| 67 | + |
| 68 | +@dataclass(frozen=True) |
| 69 | +class CheckResult: |
| 70 | + """Result of a governance check on a tool_use block or tool result.""" |
| 71 | + |
| 72 | + allowed: bool |
| 73 | + """Whether the action/result is permitted.""" |
| 74 | + |
| 75 | + block_reason: str | None = None |
| 76 | + """Reason for blocking, if not allowed.""" |
| 77 | + |
| 78 | + redacted_result: str | None = None |
| 79 | + """Redacted version of the tool result, if PII/secrets were found. |
| 80 | + Only populated by ``check_result()``, not ``check_tool_use()``.""" |
| 81 | + |
| 82 | + policies_evaluated: int = 0 |
| 83 | + """Number of policies the server evaluated.""" |
| 84 | + |
| 85 | + |
| 86 | +def _derive_connector_type(tool_name: str, action: str | None = None) -> str: |
| 87 | + """Derive AxonFlow connector_type from Computer Use tool name and action. |
| 88 | +
|
| 89 | + Examples: |
| 90 | + _derive_connector_type("computer", "left_click") -> "computer_use.left_click" |
| 91 | + _derive_connector_type("bash") -> "computer_use.bash" |
| 92 | + _derive_connector_type("text_editor") -> "computer_use.text_editor" |
| 93 | + """ |
| 94 | + if action: |
| 95 | + return f"computer_use.{action}" |
| 96 | + return f"computer_use.{tool_name}" |
| 97 | + |
| 98 | + |
| 99 | +class ComputerUseGovernor: |
| 100 | + """Governance middleware for Anthropic Computer Use tool_use blocks. |
| 101 | +
|
| 102 | + Insert into the sampling loop to evaluate every tool action against |
| 103 | + AxonFlow policies before execution, and scan results before they |
| 104 | + reach Claude's context. |
| 105 | +
|
| 106 | + Args: |
| 107 | + client: An authenticated :class:`axonflow.AxonFlow` client. |
| 108 | + blocked_bash_patterns: Regex patterns for bash commands to block |
| 109 | + client-side before even calling the server. Defaults to |
| 110 | + common destructive/exfiltration patterns. |
| 111 | + """ |
| 112 | + |
| 113 | + def __init__( |
| 114 | + self, |
| 115 | + client: AxonFlow, |
| 116 | + *, |
| 117 | + blocked_bash_patterns: list[str] | None = None, |
| 118 | + ) -> None: |
| 119 | + self._client = client |
| 120 | + self._bash_patterns = [ |
| 121 | + re.compile(p, re.IGNORECASE) |
| 122 | + for p in ( |
| 123 | + blocked_bash_patterns |
| 124 | + if blocked_bash_patterns is not None |
| 125 | + else DEFAULT_BLOCKED_BASH_PATTERNS |
| 126 | + ) |
| 127 | + ] |
| 128 | + |
| 129 | + def _check_bash_locally(self, command: str) -> str | None: |
| 130 | + """Check bash command against local blocked patterns. |
| 131 | +
|
| 132 | + Returns the matching pattern description if blocked, None if allowed. |
| 133 | + """ |
| 134 | + for pattern in self._bash_patterns: |
| 135 | + if pattern.search(command): |
| 136 | + return f"Blocked by local pattern: {pattern.pattern}" |
| 137 | + return None |
| 138 | + |
| 139 | + async def check_tool_use(self, tool_use_block: dict[str, Any]) -> CheckResult: |
| 140 | + """Check a tool_use block before execution. |
| 141 | +
|
| 142 | + Args: |
| 143 | + tool_use_block: The tool_use content block from Claude's response. |
| 144 | + Expected keys: ``name`` (str), ``input`` (dict). |
| 145 | + Optional: ``id`` (str), ``type`` (str). |
| 146 | +
|
| 147 | + Returns: |
| 148 | + CheckResult indicating whether the tool call is allowed. |
| 149 | + """ |
| 150 | + name = tool_use_block.get("name", "unknown") |
| 151 | + tool_input = tool_use_block.get("input", {}) |
| 152 | + |
| 153 | + # For bash tools (any version), check locally first (fast, no network) |
| 154 | + if name == "bash" or name.startswith("bash_"): |
| 155 | + command = tool_input.get("command", "") |
| 156 | + if isinstance(command, str): |
| 157 | + local_block = self._check_bash_locally(command) |
| 158 | + if local_block is not None: |
| 159 | + return CheckResult( |
| 160 | + allowed=False, |
| 161 | + block_reason=local_block, |
| 162 | + policies_evaluated=0, |
| 163 | + ) |
| 164 | + |
| 165 | + # Derive connector_type from tool name + action |
| 166 | + action = tool_input.get("action") if isinstance(tool_input, dict) else None |
| 167 | + connector_type = _derive_connector_type( |
| 168 | + name, |
| 169 | + action if isinstance(action, str) else None, |
| 170 | + ) |
| 171 | + |
| 172 | + # Serialize input for policy evaluation |
| 173 | + statement = json.dumps(tool_input, default=str) |
| 174 | + |
| 175 | + check = await self._client.mcp_check_input( |
| 176 | + connector_type=connector_type, |
| 177 | + statement=statement, |
| 178 | + operation="execute", |
| 179 | + ) |
| 180 | + |
| 181 | + if not check.allowed: |
| 182 | + return CheckResult( |
| 183 | + allowed=False, |
| 184 | + block_reason=check.block_reason or "Blocked by AxonFlow policy", |
| 185 | + policies_evaluated=check.policies_evaluated, |
| 186 | + ) |
| 187 | + |
| 188 | + return CheckResult( |
| 189 | + allowed=True, |
| 190 | + policies_evaluated=check.policies_evaluated, |
| 191 | + ) |
| 192 | + |
| 193 | + async def check_result( |
| 194 | + self, |
| 195 | + tool_name: str, |
| 196 | + result: str, |
| 197 | + ) -> CheckResult: |
| 198 | + """Check a tool execution result before feeding it back to Claude. |
| 199 | +
|
| 200 | + Args: |
| 201 | + tool_name: The tool name (e.g., "computer", "bash", "text_editor"). |
| 202 | + result: The text result from tool execution. |
| 203 | +
|
| 204 | + Returns: |
| 205 | + CheckResult with allowed status and optional redacted_result. |
| 206 | + """ |
| 207 | + connector_type = _derive_connector_type(tool_name) |
| 208 | + |
| 209 | + check = await self._client.mcp_check_output( |
| 210 | + connector_type=connector_type, |
| 211 | + message=result, |
| 212 | + ) |
| 213 | + |
| 214 | + if not check.allowed: |
| 215 | + return CheckResult( |
| 216 | + allowed=False, |
| 217 | + block_reason=check.block_reason or "Tool result blocked by policy", |
| 218 | + policies_evaluated=check.policies_evaluated, |
| 219 | + ) |
| 220 | + |
| 221 | + if check.redacted_data is not None: |
| 222 | + return CheckResult( |
| 223 | + allowed=True, |
| 224 | + redacted_result=( |
| 225 | + check.redacted_data |
| 226 | + if isinstance(check.redacted_data, str) |
| 227 | + else json.dumps(check.redacted_data, default=str) |
| 228 | + ), |
| 229 | + policies_evaluated=check.policies_evaluated, |
| 230 | + ) |
| 231 | + |
| 232 | + return CheckResult( |
| 233 | + allowed=True, |
| 234 | + policies_evaluated=check.policies_evaluated, |
| 235 | + ) |
0 commit comments