Coverage for merco/sandbox/guard.py: 98%

1"""工具执行守卫 — 细粒度敏感命令规则链 

2 

3每条规则 = tool + pattern + action，依次匹配，首个命中生效:: 

4 

5 guard = ToolGuard() 

6 guard.rule("bash", "DROP TABLE", "deny") # 用户可加硬拦截 

7 # 默认规则全部 ask，确认后放行 

8 

9配置 (merco.json):: 

10 

11 "sandbox_rules": [ 

12 {"tool": "bash", "pattern": "DROP TABLE", "action": "deny"}, 

13 ] 

14""" 

15 

16import logging 

17from abc import ABC, abstractmethod 

18from dataclasses import dataclass 

19from enum import Enum 

20 

21from .security import SecurityChecker 

22 

23logger = logging.getLogger("merco.guard") 

24 

25 

26# ── 枚举和结果 ────────────────────────────────────────────── 

27 

28class GuardAction(Enum): 

29 ALLOW = "allow" # 直接放行 

30 DENY = "deny" # 直接拒绝 

31 ASK = "ask" # 需要用户确认 

32 

33@dataclass 

34class GuardResult: 

35 action: GuardAction 

36 command: str 

37 rule: "GuardRule | None" = None 

38 reason: str = "" 

39 

40 

41class GuardConfirmationRequired(Exception): 

42 """需要用户确认才能继续执行""" 

43 

44 def __init__(self, result: GuardResult): 

45 self.result = result 

46 super().__init__(f"需要确认: {result.command} - {result.reason}") 

47 

48 

49# ── 规则 ────────────────────────────────────────────────── 

50 

51@dataclass 

52class GuardRule: 

53 tool: str # "bash" | "write_file" | "*" 所有工具 

54 pattern: str 

55 action: str # "ask" | "deny" | "allow" 

56 

57 def to_dict(self) -> dict: 

58 return {"tool": self.tool, "pattern": self.pattern, "action": self.action} 

59 

60 @classmethod 

61 def from_dict(cls, d: dict) -> "GuardRule": 

62 return cls( 

63 tool=d.get("tool", "bash"), 

64 pattern=d["pattern"], 

65 action=d.get("action", "ask"), 

66 ) 

67 

68 

69# ── 默认规则（全部 ask，不硬拦截）───────────────────────── 

70 

71_DEFAULT_RULES: list[GuardRule] = [ 

72 GuardRule("bash", "rm -rf /", "ask"), 

73 GuardRule("bash", "mkfs", "ask"), 

74 GuardRule("bash", "dd if=", "ask"), 

75 GuardRule("bash", "> /dev/sd", "ask"), 

76 GuardRule("bash", "curl | bash", "ask"), 

77 GuardRule("bash", "curl | sh", "ask"), 

78 GuardRule("bash", "wget | bash", "ask"), 

79 GuardRule("bash", "wget | sh", "ask"), 

80 GuardRule("bash", "chmod 777 /", "ask"), 

81 GuardRule("bash", "rm ", "ask"), 

82 GuardRule("bash", "sudo ", "ask"), 

83 GuardRule("bash", "chmod ", "ask"), 

84 GuardRule("bash", "chown ", "ask"), 

85 GuardRule("bash", "kill ", "ask"), 

86 GuardRule("bash", "pkill ", "ask"), 

87 GuardRule("bash", "git push", "ask"), 

88 GuardRule("bash", "git reset --hard", "ask"), 

89 GuardRule("bash", "docker rm", "ask"), 

90 GuardRule("bash", "docker rmi", "ask"), 

91 GuardRule("bash", "pip install", "ask"), 

92 GuardRule("bash", "pip uninstall", "ask"), 

93 GuardRule("bash", "npm install -g", "ask"), 

94 GuardRule("bash", "npm uninstall -g", "ask"), 

95 GuardRule("bash", "apt ", "ask"), 

96 GuardRule("bash", "yum ", "ask"), 

97 GuardRule("bash", "brew ", "ask"), 

98 GuardRule("bash", "shutdown", "ask"), 

99 GuardRule("bash", "reboot", "ask"), 

100] 

101 

102 

103# ── ToolGuard ───────────────────────────────────────────── 

104 

105class ToolGuard: 

106 """工具执行守卫 — facade，委托给 PolicyPipeline 

107 

108 用法:: 

109 

110 guard = ToolGuard() 

111 guard.rule("bash", "DROP TABLE", "deny") 

112 

113 result = await guard.check("bash", {"command": "rm file.txt"}) 

114 if result.action == GuardAction.DENY: 

115 return # 已拦截 

116 

117 插件用法:: 

118 

119 pipeline = PolicyPipeline() 

120 pipeline.use(BuiltinDefaultPolicy(mode="ask")) 

121 pipeline.use(MyCustomPolicy()) 

122 guard = ToolGuard(pipeline=pipeline) 

123 """ 

124 

125 def __init__(self, pipeline: "PolicyPipeline" = None, mode: str = "ask", user_rules: list = None): 

126 if pipeline: 

127 self._pipeline = pipeline 

128 else: 

129 # 向后兼容：没传 pipeline 时自动创建默认 pipeline 

130 self._pipeline = PolicyPipeline() 

131 self._pipeline.use(BuiltinDefaultPolicy(mode=mode, user_rules=user_rules)) 

132 

133 def rule(self, tool: str, pattern: str, action: str) -> "ToolGuard": 

134 """添加规则（向后兼容）。在默认策略前插入一条自定义策略。""" 

135 self._pipeline._policies.insert(0, 

136 _SingleRulePolicy(tool, pattern, action)) 

137 return self 

138 

139 async def check(self, tool_name: str, arguments: dict) -> GuardResult: 

140 return await self._pipeline.check(tool_name, arguments) 

141 

142 

143# ── PermissionPolicy ─────────────────────────────────────── 

144 

145class PermissionPolicy(ABC): 

146 """安全策略基类""" 

147 name: str = "" 

148 

149 @abstractmethod 

150 async def check(self, tool_name: str, arguments: dict) -> GuardResult | None: 

151 """检查工具。返回 GuardResult = 决断，None = 传给下一个策略""" 

152 ... 

153 

154 

155class PolicyPipeline: 

156 """安全策略责任链""" 

157 

158 def __init__(self): 

159 self._policies: list[PermissionPolicy] = [] 

160 

161 def use(self, policy: PermissionPolicy) -> "PolicyPipeline": 

162 """注册策略""" 

163 self._policies.append(policy) 

164 return self 

165 

166 async def check(self, tool_name: str, arguments: dict) -> GuardResult: 

167 """依次检查，首个返回非 None 的结果生效""" 

168 for policy in self._policies: 

169 result = await policy.check(tool_name, arguments) 

170 if result is not None: 

171 return result 

172 return GuardResult(action=GuardAction.ALLOW, command="") 

173 

174 

175class _SingleRulePolicy(PermissionPolicy): 

176 """单条规则的策略包装""" 

177 name = "single_rule" 

178 

179 def __init__(self, tool, pattern, action): 

180 self._tool = tool 

181 self._pattern = pattern 

182 self._action = action 

183 

184 async def check(self, tool_name, arguments): 

185 command = arguments.get("command", "") 

186 if self._tool not in (tool_name, "*"): 

187 return None 

188 if self._pattern not in command: 

189 return None 

190 return GuardResult(action=GuardAction(self._action), command=command) 

191 

192 

193# ── BuiltinDefaultPolicy ─────────────────────────────────── 

194 

195class BuiltinDefaultPolicy(PermissionPolicy): 

196 """默认安全策略 — 包装 30 条默认规则 + SecurityChecker + mode logic""" 

197 name = "builtin_default" 

198 

199 def __init__(self, mode: str = "ask", user_rules: list = None): 

200 self.mode = mode 

201 self._rules: list[GuardRule] = [] 

202 if user_rules: 

203 for r in user_rules: 

204 self._rules.append( 

205 GuardRule.from_dict(r) if isinstance(r, dict) else r 

206 ) 

207 self._rules.extend(_DEFAULT_RULES) 

208 

209 async def check(self, tool_name: str, arguments: dict) -> GuardResult | None: 

210 if self.mode == "auto": 

211 return GuardResult(action=GuardAction.ALLOW, command="") 

212 

213 command = arguments.get("command", "") 

214 path = arguments.get("path", "") 

215 

216 if path and tool_name != "bash": 

217 ok, reason = SecurityChecker.check_file_path(path) 

218 if not ok: 

219 return GuardResult(action=GuardAction.DENY, command=path, reason=reason) 

220 

221 if command: 

222 ok, reason = SecurityChecker.check_command(command) 

223 if not ok: 

224 return GuardResult(action=GuardAction.DENY, command=command, reason=reason) 

225 

226 for rule in self._rules: 

227 if not self._tool_match(rule.tool, tool_name): 

228 continue 

229 if rule.pattern not in command: 

230 continue 

231 action = GuardAction(rule.action) 

232 return GuardResult(action=action, command=command, rule=rule) 

233 

234 return GuardResult(action=GuardAction.ALLOW, command=command) 

235 

236 @staticmethod 

237 def _tool_match(pattern: str, tool_name: str) -> bool: 

238 return pattern == "*" or pattern == tool_name