第23章 添加工具系统与权限控制 —— Mini-Harness 的安全升级

# 权限检查模块

import re

class PermissionSystem:
    """
    权限检查系统
    """
    
    def __init__(self, config=None):
        """
        初始化权限系统
        
        参数:
            config: 权限规则配置
        """
        self.config = config or self.default_config()
    
    def default_config(self):
        """默认权限配置"""
        return {
            "deny": [
                {"tool": "Bash", "command": "rm -rf /"},
                {"tool": "Bash", "command": "rm -rf /*"},
                {"tool": "Bash", "command": "sudo"},
                {"tool": "Read", "file_path": ".env"},
                {"tool": "Read", "file_path": "*credentials*"}
            ],
            "ask": [
                {"tool": "Bash", "command": "git push"},
                {"tool": "Bash", "command": "npm install"},
                {"tool": "Bash", "command": "pip install"},
                {"tool": "Write"},
                {"tool": "Edit"}
            ],
            "allow": [
                {"tool": "Read"},
                {"tool": "Glob"},
                {"tool": "Grep"}
            ]
        }
    
    def check_permission(self, tool_name, tool_params):
        """
        检查权限
        
        参数:
            tool_name: 工具名称
            tool_params: 工具参数
        
        返回:
            permission: "deny" / "ask" / "allow"
            reason: 原因说明
        """
        # 1. 检查 deny 规则
        for rule in self.config["deny"]:
            if self.match_rule(tool_name, tool_params, rule):
                return "deny", rule
        
        # 2. 检查 ask 规则
        for rule in self.config["ask"]:
            if self.match_rule(tool_name, tool_params, rule):
                return "ask", rule
        
        # 3. 检查 allow 规则
        for rule in self.config["allow"]:
            if self.match_rule(tool_name, tool_params, rule):
                return "allow", None
        
        # 4. 无匹配时默认 ask（保守策略）
        return "ask", {"reason": "默认需要审批"}
    
    def match_rule(self, tool_name, tool_params, rule):
        """
        匹配规则
        
        参数:
            tool_name: 工具名称
            tool_params: 工具参数
            rule: 权限规则
        
        返回:
            matched: 是否匹配
        """
        # 1. 工具名称匹配
        if rule.get("tool") and rule["tool"] != tool_name:
            return False
        
        # 2. 命令匹配（Bash 工具）
        if rule.get("command") and tool_name == "Bash":
            command = tool_params.get("command", "")
            # 模式匹配（支持通配符）
            if not self.pattern_match(command, rule["command"]):
                return False
        
        # 3. 文件路径匹配（Read/Write/Edit）
        if rule.get("file_path") and tool_name in ["Read", "Write", "Edit"]:
            file_path = tool_params.get("file_path", "")
            # 模式匹配（支持通配符）
            if not self.pattern_match(file_path, rule["file_path"]):
                return False
        
        return True
    
    def pattern_match(self, value, pattern):
        """
        模式匹配（支持通配符）
        
        参数:
            value: 实际值
            pattern: 匹配模式
        
        返回:
            matched: 是否匹配
        """
        # 转换通配符为正则
        regex_pattern = pattern.replace("*", ".*")
        regex_pattern = regex_pattern.replace("?", ".")
        
        # 正则匹配
        return re.match(regex_pattern, value) is not None


# 测试权限检查
def test_permission_system():
    perm = PermissionSystem()
    
    # 测试 deny
    result = perm.check_permission("Bash", {"command": "rm -rf /"})
    print(f"rm -rf /: {result}")
    
    # 测试 ask
    result = perm.check_permission("Bash", {"command": "git push"})
    print(f"git push: {result}")
    
    # 测试 allow
    result = perm.check_permission("Read", {"file_path": "app.ts"})
    print(f"Read app.ts: {result}")
    
    # 测试默认 ask
    result = perm.check_permission("Bash", {"command": "ls -la"})
    print(f"ls -la: {result}")


# 运行测试
test_permission_system()

# 用户审批模块

class ApprovalSystem:
    """
    用户审批系统
    """
    
    def __init__(self):
        self.pending_requests = []
    
    def request_approval(self, tool_name, tool_params, reason):
        """
        请求用户审批
        
        参数:
            tool_name: 工具名称
            tool_params: 工具参数
            reason: 原因
        
        返回:
            approved: 是否批准
        """
        # 显示审批提示
        prompt = self.format_approval_prompt(tool_name, tool_params, reason)
        
        print(prompt)
        
        # 等待用户输入
        response = input("是否批准? (y/n): ").strip().lower()
        
        approved = response == 'y'
        
        return approved
    
    def format_approval_prompt(self, tool_name, tool_params, reason):
        """
        格式化审批提示
        
        参数:
            tool_name: 工具名称
            tool_params: 工具参数
            reason: 原因
        
        返回:
            prompt: 审批提示文本
        """
        reason_text = reason.get("reason", reason.get("command", "需要审批"))
        
        prompt = f"""
⚠️ 权限审批请求

工具: {tool_name}
参数: {self.format_params(tool_params)}
原因: {reason_text}

是否批准执行？
[y] 批准  [n] 拒绝
"""
        
        return prompt
    
    def format_params(self, params):
        """格式化参数"""
        if not params:
            return "无参数"
        
        formatted = []
        for key, value in params.items():
            if len(str(value)) > 50:
                value = str(value)[:50] + "..."
            formatted.append(f"{key}: {value}")
        
        return "\n    " + "\n    ".join(formatted)


# 测试审批系统
def test_approval_system():
    approval = ApprovalSystem()
    
    # 模拟请求审批
    approved = approval.request_approval(
        "Bash",
        {"command": "git push origin main"},
        {"tool": "Bash", "command": "git push"}
    )
    
    print(f"用户响应: {approved}")


# 运行测试（实际交互）
# test_approval_system()

# Mini-Harness + 权限控制完整实现

import os
import subprocess
import re

class MiniHarnessWithPermissions:
    """
    Mini-Harness + 权限控制
    """
    
    def __init__(self, permission_config=None):
        """
        初始化
        
        参数:
            permission_config: 权限配置
        """
        self.system_prompt = self.default_system_prompt()
        self.tools = self.create_tools()
        self.permission_system = PermissionSystem(permission_config)
        self.approval_system = ApprovalSystem()
        self.state = None
    
    def default_system_prompt(self):
        """默认 System Prompt"""
        return """
你是代码助手 Agent。你可以使用以下工具：

工具列表：
- Read: 读取文件（参数：file_path）
- Write: 写入文件（参数：file_path, content）
- Edit: 编辑文件（参数：file_path, old_string, new_string）
- Bash: 执行命令（参数：command）
- Glob: 搜索文件（参数：pattern）
- Grep: 搜索内容（参数：pattern, path）

约束：
- 每次只调用一个工具
- 工具调用后等待结果
- 完成任务后返回 end_turn
"""
    
    def create_tools(self):
        """创建工具映射"""
        return {
            "Read": self.tool_read,
            "Write": self.tool_write,
            "Edit": self.tool_edit,
            "Bash": self.tool_bash,
            "Glob": self.tool_glob,
            "Grep": self.tool_grep
        }
    
    # 工具实现（略，与上一篇相同）
    
    def tool_read(self, params):
        """Read 工具"""
        file_path = params.get("file_path")
        if not file_path:
            return "错误：缺少 file_path 参数"
        if not os.path.exists(file_path):
            return f"错误：文件 {file_path} 不存在"
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            return f"文件内容:\n{content}"
        except Exception as e:
            return f"错误：读取文件失败 - {str(e)}"
    
    def tool_write(self, params):
        """Write 工具"""
        file_path = params.get("file_path")
        content = params.get("content")
        if not file_path or not content:
            return "错误：缺少 file_path 或 content 参数"
        try:
            os.makedirs(os.path.dirname(file_path), exist_ok=True)
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(content)
            return f"成功：文件 {file_path} 已写入"
        except Exception as e:
            return f"错误：写入文件失败 - {str(e)}"
    
    def tool_edit(self, params):
        """Edit 工具"""
        file_path = params.get("file_path")
        old_string = params.get("old_string")
        new_string = params.get("new_string")
        
        if not file_path or not old_string or not new_string:
            return "错误：缺少必要参数"
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            if old_string not in content:
                return "错误：old_string 未找到"
            
            # 替换（简单实现，只替换第一个匹配）
            new_content = content.replace(old_string, new_string, 1)
            
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(new_content)
            
            return f"成功：文件 {file_path} 已编辑"
        except Exception as e:
            return f"错误：编辑文件失败 - {str(e)}"
    
    def tool_bash(self, params):
        """Bash 工具"""
        command = params.get("command")
        if not command:
            return "错误：缺少 command 参数"
        try:
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=30
            )
            if result.returncode == 0:
                return f"命令输出:\n{result.stdout}"
            else:
                return f"错误：命令执行失败\n{result.stderr}"
        except subprocess.TimeoutExpired:
            return "错误：命令执行超时（30秒）"
        except Exception as e:
            return f"错误：执行命令失败 - {str(e)}"
    
    def tool_glob(self, params):
        """Glob 工具"""
        pattern = params.get("pattern")
        if not pattern:
            return "错误：缺少 pattern 参数"
        import glob
        files = glob.glob(pattern, recursive=True)
        return f"匹配文件:\n{chr(10).join(files)}"
    
    def tool_grep(self, params):
        """Grep 工具"""
        pattern = params.get("pattern")
        path = params.get("path", ".")
        if not pattern:
            return "错误：缺少 pattern 参数"
        try:
            result = subprocess.run(
                f"grep -r {pattern} {path}",
                shell=True,
                capture_output=True,
                text=True
            )
            return f"搜索结果:\n{result.stdout}"
        except Exception as e:
            return f"错误：搜索失败 - {str(e)}"
    
    def initialize(self, prompt):
        """初始化"""
        self.state = {
            "messages": [
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": prompt}
            ],
            "loops": 0,
            "tool_calls": 0,
            "denied": 0,
            "approved": 0
        }
        return self.state
    
    def agent_loop(self, model_api):
        """Agent Loop + 权限检查"""
        while True:
            messages = self.state["messages"]
            
            # Model 推理
            response = model_api(messages)
            
            stop_reason = response.get("stop_reason", "end_turn")
            
            if stop_reason == "end_turn":
                return response
            
            elif stop_reason == "tool_use":
                tool_use = self.extract_tool_use(response)
                
                tool_name = tool_use["name"]
                tool_params = tool_use["parameters"]
                
                # 权限检查
                permission, reason = self.permission_system.check_permission(
                    tool_name, tool_params
                )
                
                if permission == "deny":
                    # 禁止执行
                    tool_result = f"权限拒绝：{reason}"
                    self.state["denied"] += 1
                
                elif permission == "ask":
                    # 请求审批
                    approved = self.approval_system.request_approval(
                        tool_name, tool_params, reason
                    )
                    
                    if approved:
                        # 执行工具
                        tool_result = self.tools[tool_name](tool_params)
                        self.state["approved"] += 1
                    else:
                        # 拒绝执行
                        tool_result = "用户拒绝执行"
                        self.state["denied"] += 1
                
                elif permission == "allow":
                    # 自动执行
                    tool_result = self.tools[tool_name](tool_params)
                
                else:
                    # 默认 ask
                    approved = self.approval_system.request_approval(
                        tool_name, tool_params, {"reason": "默认需要审批"}
                    )
                    
                    if approved:
                        tool_result = self.tools[tool_name](tool_params)
                    else:
                        tool_result = "用户拒绝执行"
                
                # 追加结果
                self.state["messages"].append({
                    "role": "assistant",
                    "content": response["content"]
                })
                
                self.state["messages"].append({
                    "role": "user",
                    "content": tool_result
                })
                
                # 更新状态
                self.state["loops"] += 1
                self.state["tool_calls"] += 1
                
                # 循环限制
                if self.state["loops"] > 20:
                    print("警告：循环超过 20 次，强制结束")
                    return response
    
    def extract_tool_use(self, response):
        """提取 tool_use"""
        content = response.get("content", "")
        
        import json
        
        json_match = re.search(r'\{[^}]+\}', content)
        if json_match:
            tool_info = json.loads(json_match.group())
            return {
                "name": tool_info.get("tool", "unknown"),
                "parameters": tool_info.get("params", {})
            }
        
        return {"name": "unknown", "parameters": {}}
    
    def run(self, prompt, model_api):
        """完整运行"""
        self.initialize(prompt)
        response = self.agent_loop(model_api)
        
        return {
            "response": response,
            "loops": self.state["loops"],
            "tool_calls": self.state["tool_calls"],
            "denied": self.state["denied"],
            "approved": self.state["approved"],
            "messages": self.state["messages"]
        }


# PermissionSystem 和 ApprovalSystem 类（前面已定义）

# 测试 Mini-Harness + 权限控制
def test_mini_harness_with_permissions():
    harness = MiniHarnessWithPermissions()
    
    prompt = "执行 git push 命令"
    
    # 模拟 model_api（返回 git push 工具调用）
    def mock_model_api(messages):
        return {
            "content": "决定执行 git push。{\"tool\": \"Bash\", \"params\": {\"command\": \"git push\"}}",
            "stop_reason": "tool_use"
        }
    
    result = harness.run(prompt, mock_model_api)
    
    print("运行完成:")
    print(f"- 循环次数: {result['loops']}")
    print(f"- 工具调用次数: {result['tool_calls']}")
    print(f"- 拒绝次数: {result['denied']}")
    print(f"- 批准次数: {result['approved']}")


# 运行测试
test_mini_harness_with_permissions()