深度驾驭 Anthropic Agent SDK 钩子：通过 pretooluse 与 posttooluse 构建安全闭环

import re
from typing import Dict, Any, Tuple
from anthropic import Anthropic
from anthropic.types.beta import BetaToolUseBlock

# 定义危险 Shell 命令黑名单正则
PROHIBITED_COMMAND_PATTERN = re.compile(
    r"\b(rm\s+-rf|mkfs|dd|chmod|chown|shutdown|reboot|passwd|wget|curl)\b",
    re.IGNORECASE
)

# 敏感 Key 与凭证匹配模式
SECRET_PATTERN = re.compile(r"(api[-_]?key|passwd|password|secret|token)[\s=:\"]+([A-Za-z0-9_\-\.]{12,})", re.IGNORECASE)

class SecurityAuditException(Exception):
    """安全审计不通过时抛出的异常"""
    pass

def pre_tool_use_hook(tool_name: str, args: Dict[str, Any]) -> Tuple[bool, str]:
    """
    在工具执行前进行入参拦截审计。
    返回 (是否批准, 错误提示或修改后的参数)
    """
    print(f"[PRE-TOOL-USE] 正在审计工具: {tool_name}")
    
    if tool_name == "execute_bash":
        command = args.get("command", "")
        
        # 1. 拦截高危的系统指令和删除操作
        if PROHIBITED_COMMAND_PATTERN.search(command):
            print(f"[安全警报] 成功拦截破坏性 Bash 执行意图: {command}")
            return False, "Error: The command you attempted to run is prohibited by security policy."
            
        # 2. 拦截对系统敏感路径的越权访问
        if "/etc/" in command or "/var/log" in command:
            print(f"[安全警报] 成功拦截敏感路径越权访问: {command}")
            return False, "Error: Access to system configuration and log directories is restricted."
            
        # 3. 拦截尝试在命令中硬编码写出凭证的行为（防御凭证外泄）
        if SECRET_PATTERN.search(command):
            print(f"[安全警报] 检测到高危行为：命令中包含硬编码的敏感密钥。")
            return False, "Error: Command aborted. Writing plain-text credentials is prohibited."

    return True, ""

def post_tool_use_hook(tool_name: str, raw_output: str, max_char_limit: int = 5000) -> str:
    """
    在工具执行完毕、结果返回给大模型前进行拦截。
    - 对输出内容中的密钥进行敏感信息脱敏
    - 对超长内容进行智能截断，并附加结构化说明，保护大模型上下文
    """
    print(f"[POST-TOOL-USE] 拦截并处理工具输出结果: {tool_name}")
    
    # 1. 抹除输出中的密钥敏感数据（例如报错信息带出来的 DSN）
    sanitized_output = SECRET_PATTERN.sub(r"\1=REDACTED_SECRET", raw_output)
    
    # 2. 强行截断超大 Payload，实现 Token 成本控制
    if len(sanitized_output) > max_char_limit:
        truncated_size = len(sanitized_output) - max_char_limit
        print(f"[性能预警] 工具输出体积过大 ({len(sanitized_output)} 字符). 启动截断程序。")
        
        # 截断超长部分，但在末尾附加极其清晰的“结构化声明”
        # 告诉 Claude 这里被截断了，并引导它通过更精确的参数（如 Limit/Filter）进行二次查询。
        truncated_part = sanitized_output[:max_char_limit]
        optimized_output = (
            f"{truncated_part}\n\n"
            f"[SYSTEM WARNING: Output truncated. Kept first {max_char_limit} characters. "
            f"Omitted remaining {truncated_size} characters to save context tokens. "
            f"If you need to query a specific section, refine your tool parameters with a selective filter or limit.]"
        )
        return optimized_output
        
    return sanitized_output

import json
from anthropic import Anthropic

client = Anthropic()

def execute_agent_turn(task: str):
    messages = [{"role": "user", "content": task}]
    tools = [
        {
            "name": "execute_bash",
            "description": "Executes shell commands on the terminal.",
            "input_schema": {
                "type": "object",
                "properties": {
                    "command": {"type": "string", "description": "The command to run."}
                },
                "required": ["command"]
            }
        }
    ]
    
    # 显式设定最大 Turn 次数，防止 Agent 在死循环中狂烧额度
    for turn in range(5):  
        response = client.beta.messages.create(
            model="claude-3-5-sonnet-latest",
            max_tokens=2048,
            tools=tools,
            messages=messages
        )
        
        # 记录大模型这一轮的决策
        messages.append({"role": "assistant", "content": response.content})
        
        tool_calls = [block for block in response.content if block.type == "tool_use"]
        
        # 如果大模型不需要调用工具，说明任务已收敛
        if not tool_calls:
            print("[Agent 执行收敛] 任务圆满完成。")
            print(f"最终结论: {response.content[0].text}")
            break
            
        tool_results = []
        for tool_use in tool_calls:
            tool_name = tool_use.name
            tool_args = tool_use.input
            tool_id = tool_use.id
            
            # --- 1. 执行 pretooluse 入参安全审计钩子 ---
            is_approved, check_result = pre_tool_use_hook(tool_name, tool_args)
            
            if not is_approved:
                # 审计失败，安全断路，直接强行将错误返回给 Claude
                tool_results.append({
                    "role": "user",
                    "content": [
                        {
                            "type": "tool_result",
                            "tool_use_id": tool_id,
                            "content": check_result,
                            "is_error": True
                        }
                    ]
                })
                continue
                
            # --- 2. 工具实际运行（此处使用 Mock 逻辑示意） ---
            try:
                if tool_name == "execute_bash":
                    raw_result = f"Successfully executed command: {tool_args['command']}. Returned 15000 records..."
                else:
                    raw_result = "Execution error"
            except Exception as e:
                raw_result = f"Error executing tool: {str(e)}"
                
            # --- 3. 执行 posttooluse 出参优化与脱敏钩子 ---
            optimized_result = post_tool_use_hook(tool_name, raw_result)
            
            tool_results.append({
                "role": "user",
                "content": [
                    {
                        "type": "tool_result",
                        "tool_use_id": tool_id,
                        "content": optimized_result
                    }
                ]
            })
            
        # 将本轮工具运行的反馈数据追加进上下文，准备开启下一轮大模型 Turn
        messages.extend(tool_results)