Coverage for agentos/system/shell_exec.py: 33%

1""" 

2Shell 执行模块 — 带权限检查和安全沙箱的命令执行。 

3 

4分层策略: 

5- SHELL_READONLY: 只允许预定义安全命令白名单 

6- SHELL_STANDARD: 允许任意命令，超时+沙箱目录限制 

7- SHELL_FULL: 无限制（需二次确认） 

8""" 

9 

10from __future__ import annotations 

11 

12import os 

13import re 

14import signal 

15import subprocess 

16import tempfile 

17import shlex 

18from dataclasses import dataclass, field 

19 

20from agentos.system.permissions import ( 

21 SystemPermissionManager, 

22 PermissionTier, 

23 PermissionDenied, 

24) 

25 

26 

27# ── Shell 策略 ───────────────────────────────────────────────── 

28 

29 

30@dataclass 

31class ShellPolicy: 

32 """Shell 执行策略。""" 

33 allowed_commands: list[str] = field(default_factory=list) # 命令白名单 

34 blocked_commands: list[str] = field(default_factory=list) # 命令黑名单 

35 blocked_patterns: list[str] = field(default_factory=list) # 参数黑名单模式 

36 timeout_seconds: int = 30 # 超时时间 

37 max_output_bytes: int = 1024 * 100 # 最大输出字节 

38 sandbox_dir: str = "" # 沙箱工作目录 

39 allow_pipes: bool = False # 是否允许管道 

40 allow_redirects: bool = False # 是否允许重定向 

41 

42 

43@dataclass 

44class ShellResult: 

45 """Shell 执行结果。""" 

46 success: bool 

47 command: str 

48 stdout: str 

49 stderr: str 

50 exit_code: int = -1 

51 duration_ms: float = 0 

52 timeout: bool = False 

53 permission_denied: bool = False 

54 error: str = "" 

55 

56 

57# ── 预设策略 ─────────────────────────────────────────────────── 

58 

59READONLY_POLICY = ShellPolicy( 

60 allowed_commands=[ 

61 "ls", "cat", "head", "tail", "find", "ps", "df", "du", 

62 "whoami", "pwd", "env", "echo", "date", "wc", "stat", 

63 "file", "which", "uname", "uptime", "id", "groups", 

64 "free", "top", "grep", "awk", "sed", "sort", "uniq", 

65 "cut", "tr", "tee", "xargs", "basename", "dirname", 

66 "readlink", "realpath", "md5sum", "sha256sum", "diff", 

67 "curl", "wget", "ping", "hostname", "ip", "ss", 

68 "python3", "python", "pip", "pip3", "git", "node", "npm", 

69 ], 

70 timeout_seconds=30, 

71 allow_pipes=True, 

72 allow_redirects=False, 

73) 

74 

75STANDARD_POLICY = ShellPolicy( 

76 blocked_commands=[ 

77 "rm", "shutdown", "reboot", "halt", "poweroff", 

78 "mkfs", "dd", "fdisk", "parted", "mount", "umount", 

79 "chmod", "chown", "useradd", "userdel", "passwd", 

80 "iptables", "ufw", "systemctl", "service", 

81 ], 

82 blocked_patterns=[ 

83 r"rm\s+(-rf?|--recursive)\s+/", # rm -rf / 

84 r">\s*/dev/", # 覆盖设备 

85 r"mkfs\.", # 格式化 

86 r"dd\s+if=", # dd 操作 

87 r"curl.*\|.*sh", # curl pipe sh 

88 r"wget.*\|.*sh", # wget pipe sh 

89 r">\s*/etc/", # 写入 /etc/ 

90 ], 

91 timeout_seconds=60, 

92 max_output_bytes=1024 * 500, 

93 allow_pipes=True, 

94 allow_redirects=True, 

95) 

96 

97FULL_POLICY = ShellPolicy( 

98 timeout_seconds=300, 

99 max_output_bytes=1024 * 1024 * 10, 

100 allow_pipes=True, 

101 allow_redirects=True, 

102) 

103 

104 

105# ── Shell 沙箱 ───────────────────────────────────────────────── 

106 

107 

108class ShellSandbox: 

109 """Shell 沙箱 — 隔离命令执行环境。 

110 

111 特性: 

112 - 临时工作目录隔离 

113 - 环境变量过滤（移除敏感变量） 

114 - 资源限制（超时、输出大小） 

115 - 进程组管理（确保超时时子进程也被杀死） 

116 """ 

117 

118 def __init__(self, work_dir: str | None = None): 

119 self._work_dir = work_dir or tempfile.mkdtemp(prefix="agentos_shell_") 

120 

121 @property 

122 def work_dir(self) -> str: 

123 return self._work_dir 

124 

125 def filtered_env(self) -> dict[str, str]: 

126 """返回过滤后的环境变量（移除敏感变量）。""" 

127 blocked = {"AWS_", "SECRET_", "TOKEN", "PASSWORD", "PASSWD", 

128 "KEY", "CREDENTIAL", "PRIVATE", "CERT", "AUTH"} 

129 env = {} 

130 for k, v in os.environ.items(): 

131 if not any(b in k.upper() for b in blocked): 

132 env[k] = v 

133 # 设置安全默认值 

134 env["HOME"] = self._work_dir 

135 env["PATH"] = os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin") 

136 return env 

137 

138 def cleanup(self) -> None: 

139 """清理沙箱目录。""" 

140 import shutil 

141 try: 

142 shutil.rmtree(self._work_dir, ignore_errors=True) 

143 except Exception: 

144 pass 

145 

146 

147# ── Shell 执行器 ──────────────────────────────────────────────── 

148 

149 

150class ShellExecutor: 

151 """Shell 执行器 — 带策略和权限检查的命令执行。""" 

152 

153 def __init__(self, perm_manager: SystemPermissionManager, session_id: str): 

154 self._pm = perm_manager 

155 self._sid = session_id 

156 self._sandboxes: dict[str, ShellSandbox] = {} 

157 

158 # ── 沙箱管理 ── 

159 

160 def create_sandbox(self, name: str = "default") -> ShellSandbox: 

161 """创建命名沙箱。""" 

162 sb = ShellSandbox() 

163 self._sandboxes[name] = sb 

164 return sb 

165 

166 def get_sandbox(self, name: str = "default") -> ShellSandbox: 

167 """获取或创建沙箱。""" 

168 if name not in self._sandboxes: 

169 return self.create_sandbox(name) 

170 return self._sandboxes[name] 

171 

172 def cleanup_sandbox(self, name: str = "default") -> None: 

173 """清理指定沙箱。""" 

174 sb = self._sandboxes.pop(name, None) 

175 if sb: 

176 sb.cleanup() 

177 

178 def cleanup_all(self) -> None: 

179 for sb in list(self._sandboxes.values()): 

180 sb.cleanup() 

181 self._sandboxes.clear() 

182 

183 # ── 命令执行 ── 

184 

185 def execute(self, command: str, sandbox_name: str = "default") -> ShellResult: 

186 """执行 Shell 命令，自动选择策略。""" 

187 # 选择策略 

188 try: 

189 self._pm.require(self._sid, PermissionTier.SHELL_FULL, command) 

190 policy = FULL_POLICY 

191 tier = PermissionTier.SHELL_FULL 

192 except PermissionDenied: 

193 try: 

194 self._pm.require(self._sid, PermissionTier.SHELL_STANDARD, command) 

195 policy = STANDARD_POLICY 

196 tier = PermissionTier.SHELL_STANDARD 

197 except PermissionDenied: 

198 try: 

199 self._pm.require(self._sid, PermissionTier.SHELL_READONLY, command) 

200 policy = READONLY_POLICY 

201 tier = PermissionTier.SHELL_READONLY 

202 except PermissionDenied as e: 

203 return ShellResult( 

204 success=False, command=command, 

205 stdout="", stderr="", permission_denied=True, error=str(e), 

206 ) 

207 

208 return self._execute_with_policy(command, policy, sandbox_name) 

209 

210 def execute_checked(self, command: str, required_tier: PermissionTier, 

211 sandbox_name: str = "default") -> ShellResult: 

212 """以指定权限级别执行命令。""" 

213 self._pm.require(self._sid, required_tier, command) 

214 if required_tier == PermissionTier.SHELL_READONLY: 

215 policy = READONLY_POLICY 

216 elif required_tier == PermissionTier.SHELL_STANDARD: 

217 policy = STANDARD_POLICY 

218 else: 

219 policy = FULL_POLICY 

220 return self._execute_with_policy(command, policy, sandbox_name) 

221 

222 # ── 内部实现 ── 

223 

224 def _execute_with_policy(self, command: str, policy: ShellPolicy, 

225 sandbox_name: str) -> ShellResult: 

226 """按策略执行命令。""" 

227 import time 

228 

229 # 安全检查 

230 safety_check = self._safety_check(command, policy) 

231 if safety_check: 

232 return ShellResult( 

233 success=False, command=command, 

234 stdout="", stderr=safety_check, error=safety_check, 

235 ) 

236 

237 sandbox = self.get_sandbox(sandbox_name) 

238 

239 try: 

240 t0 = time.time() 

241 proc = subprocess.Popen( 

242 command, 

243 shell=True, 

244 stdout=subprocess.PIPE, 

245 stderr=subprocess.PIPE, 

246 cwd=sandbox.work_dir, 

247 env=sandbox.filtered_env(), 

248 preexec_fn=os.setsid, # 创建新进程组，便于超时杀子进程 

249 text=True, 

250 ) 

251 

252 try: 

253 stdout, stderr = proc.communicate(timeout=policy.timeout_seconds) 

254 exit_code = proc.returncode 

255 timeout = False 

256 except subprocess.TimeoutExpired: 

257 # 杀死整个进程组 

258 os.killpg(os.getpgid(proc.pid), signal.SIGTERM) 

259 try: 

260 stdout, stderr = proc.communicate(timeout=5) 

261 except subprocess.TimeoutExpired: 

262 os.killpg(os.getpgid(proc.pid), signal.SIGKILL) 

263 stdout, stderr = proc.communicate() 

264 exit_code = -1 

265 timeout = True 

266 

267 duration_ms = (time.time() - t0) * 1000 

268 

269 # 截断过大的输出 

270 stdout = self._truncate(stdout, policy.max_output_bytes) 

271 stderr = self._truncate(stderr, policy.max_output_bytes) 

272 

273 return ShellResult( 

274 success=(exit_code == 0 and not timeout), 

275 command=command, 

276 stdout=stdout, 

277 stderr=stderr, 

278 exit_code=exit_code, 

279 duration_ms=duration_ms, 

280 timeout=timeout, 

281 ) 

282 

283 except Exception as e: 

284 return ShellResult( 

285 success=False, command=command, 

286 stdout="", stderr="", error=str(e), 

287 ) 

288 

289 def _safety_check(self, command: str, policy: ShellPolicy) -> str: 

290 """安全检查，返回错误信息或空字符串。""" 

291 # 提取主命令 

292 cmd_parts = shlex.split(command) if command else [] 

293 if not cmd_parts: 

294 return "空命令" 

295 

296 main_cmd = os.path.basename(cmd_parts[0]) 

297 

298 # 白名单检查 

299 if policy.allowed_commands: 

300 if main_cmd not in policy.allowed_commands and cmd_parts[0] not in policy.allowed_commands: 

301 return f"命令 '{main_cmd}' 不在允许列表中。允许的命令: {', '.join(policy.allowed_commands[:20])}" 

302 

303 # 黑名单检查 

304 if policy.blocked_commands: 

305 if main_cmd in policy.blocked_commands or cmd_parts[0] in policy.blocked_commands: 

306 return f"命令 '{main_cmd}' 已被阻止。被阻止的命令: {', '.join(policy.blocked_commands)}" 

307 

308 # 危险模式检查 

309 for pattern in policy.blocked_patterns: 

310 if re.search(pattern, command): 

311 return f"命令包含危险模式: {pattern}" 

312 

313 # 管道检查 

314 if not policy.allow_pipes and "|" in command: 

315 return "管道操作不被允许" 

316 

317 # 重定向检查 

318 if not policy.allow_redirects and re.search(r"[<>]", command): 

319 return "重定向操作不被允许" 

320 

321 return "" 

322 

323 @staticmethod 

324 def _truncate(text: str, max_bytes: int) -> str: 

325 """截断文本到指定字节数。""" 

326 encoded = text.encode("utf-8") 

327 if len(encoded) <= max_bytes: 

328 return text 

329 truncated = encoded[:max_bytes].decode("utf-8", errors="replace") 

330 return truncated + f"\n... [截断: {len(encoded)} → {max_bytes} 字节]"