Coverage for agentos/system/shell

1"""

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

4分层策略:

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

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

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

8"""

10from __future__ import annotations

12import os

13import re

14import shlex

15import signal

16import subprocess

17import tempfile

18from dataclasses import dataclass, field

20from agentos.system.permissions import (

21 PermissionDenied,

22 PermissionTier,

23 SystemPermissionManager,

24)

26# ── Shell 策略 ─────────────────────────────────────────────────

29@dataclass

30class ShellPolicy:

31 """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 # 是否允许重定向

43@dataclass

44class ShellResult:

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

47 success: bool

48 command: str

49 stdout: str

50 stderr: str

51 exit_code: int = -1

52 duration_ms: float = 0

53 timeout: bool = False

54 permission_denied: bool = False

55 error: str = ""

58# ── 预设策略 ───────────────────────────────────────────────────

60READONLY_POLICY = ShellPolicy(

61 allowed_commands=[

62 "ls",

63 "cat",

64 "head",

65 "tail",

66 "find",

67 "ps",

68 "df",

69 "du",

70 "whoami",

71 "pwd",

72 "env",

73 "echo",

74 "date",

75 "wc",

76 "stat",

77 "file",

78 "which",

79 "uname",

80 "uptime",

81 "id",

82 "groups",

83 "free",

84 "top",

85 "grep",

86 "awk",

87 "sed",

88 "sort",

89 "uniq",

90 "cut",

91 "tr",

92 "tee",

93 "xargs",

94 "basename",

95 "dirname",

96 "readlink",

97 "realpath",

98 "md5sum",

99 "sha256sum",

100 "diff",

101 "curl",

102 "wget",

103 "ping",

104 "hostname",

105 "ip",

106 "ss",

107 "python3",

108 "python",

109 "pip",

110 "pip3",

111 "git",

112 "node",

113 "npm",

114 ],

115 timeout_seconds=30,

116 allow_pipes=True,

117 allow_redirects=False,

118)

119

120STANDARD_POLICY = ShellPolicy(

121 blocked_commands=[

122 "rm",

123 "shutdown",

124 "reboot",

125 "halt",

126 "poweroff",

127 "mkfs",

128 "dd",

129 "fdisk",

130 "parted",

131 "mount",

132 "umount",

133 "chmod",

134 "chown",

135 "useradd",

136 "userdel",

137 "passwd",

138 "iptables",

139 "ufw",

140 "systemctl",

141 "service",

142 ],

143 blocked_patterns=[

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

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

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

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

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

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

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

151 ],

152 timeout_seconds=60,

153 max_output_bytes=1024 * 500,

154 allow_pipes=True,

155 allow_redirects=True,

156)

157

158FULL_POLICY = ShellPolicy(

159 timeout_seconds=300,

160 max_output_bytes=1024 * 1024 * 10,

161 allow_pipes=True,

162 allow_redirects=True,

163)

164

165

166# ── Shell 沙箱 ─────────────────────────────────────────────────

167

168

169class ShellSandbox:

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

171

172 特性:

173 - 临时工作目录隔离

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

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

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

177 """

178

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

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

181

182 @property

183 def work_dir(self) -> str:

184 return self._work_dir

185

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

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

188 blocked = {

189 "AWS_",

190 "SECRET_",

191 "TOKEN",

192 "PASSWORD",

193 "PASSWD",

194 "KEY",

195 "CREDENTIAL",

196 "PRIVATE",

197 "CERT",

198 "AUTH",

199 }

200 env = {}

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

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

203 env[k] = v

204 # 设置安全默认值

205 env["HOME"] = self._work_dir

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

207 return env

208

209 def cleanup(self) -> None:

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

211 import shutil

212

213 try:

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

215 except Exception:

216 pass

217

218

219# ── Shell 执行器 ────────────────────────────────────────────────

220

221

222class ShellExecutor:

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

224

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

226 self._pm = perm_manager

227 self._sid = session_id

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

229

230 # ── 沙箱管理 ──

231

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

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

234 sb = ShellSandbox()

235 self._sandboxes[name] = sb

236 return sb

237

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

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

240 if name not in self._sandboxes:

241 return self.create_sandbox(name)

242 return self._sandboxes[name]

243

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

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

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

247 if sb:

248 sb.cleanup()

249

250 def cleanup_all(self) -> None:

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

252 sb.cleanup()

253 self._sandboxes.clear()

254

255 # ── 命令执行 ──

256

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

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

259 # 选择策略

260 try:

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

262 policy = FULL_POLICY

263 except PermissionDenied:

264 try:

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

266 policy = STANDARD_POLICY

267 except PermissionDenied:

268 try:

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

270 policy = READONLY_POLICY

271 except PermissionDenied as e:

272 return ShellResult(

273 success=False,

274 command=command,

275 stdout="",

276 stderr="",

277 permission_denied=True,

278 error=str(e),

279 )

280

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

282

283 def execute_checked(

284 self, command: str, required_tier: PermissionTier, sandbox_name: str = "default"

285 ) -> ShellResult:

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

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

288 if required_tier == PermissionTier.SHELL_READONLY:

289 policy = READONLY_POLICY

290 elif required_tier == PermissionTier.SHELL_STANDARD:

291 policy = STANDARD_POLICY

292 else:

293 policy = FULL_POLICY

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

295

296 # ── 内部实现 ──

297

298 def _execute_with_policy(

299 self, command: str, policy: ShellPolicy, sandbox_name: str

300 ) -> ShellResult:

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

302 import time

303

304 # 安全检查

305 safety_check = self._safety_check(command, policy)

306 if safety_check:

307 return ShellResult(

308 success=False,

309 command=command,

310 stdout="",

311 stderr=safety_check,

312 error=safety_check,

313 )

314

315 sandbox = self.get_sandbox(sandbox_name)

316

317 try:

318 t0 = time.time()

319 proc = subprocess.Popen(

320 command,

321 shell=True,

322 stdout=subprocess.PIPE,

323 stderr=subprocess.PIPE,

324 cwd=sandbox.work_dir,

325 env=sandbox.filtered_env(),

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

327 text=True,

328 )

329

330 try:

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

332 exit_code = proc.returncode

333 timeout = False

334 except subprocess.TimeoutExpired:

335 # 杀死整个进程组

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

337 try:

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

339 except subprocess.TimeoutExpired:

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

341 stdout, stderr = proc.communicate()

342 exit_code = -1

343 timeout = True

344

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

346

347 # 截断过大的输出

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

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

350

351 return ShellResult(

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

353 command=command,

354 stdout=stdout,

355 stderr=stderr,

356 exit_code=exit_code,

357 duration_ms=duration_ms,

358 timeout=timeout,

359 )

360

361 except Exception as e:

362 return ShellResult(

363 success=False,

364 command=command,

365 stdout="",

366 stderr="",

367 error=str(e),

368 )

369

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

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

372 # 提取主命令

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

374 if not cmd_parts:

375 return "空命令"

376

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

378

379 # 白名单检查

380 if policy.allowed_commands:

381 if (

382 main_cmd not in policy.allowed_commands

383 and cmd_parts[0] not in policy.allowed_commands

384 ):

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

386

387 # 黑名单检查

388 if policy.blocked_commands:

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

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

391

392 # 危险模式检查

393 for pattern in policy.blocked_patterns:

394 if re.search(pattern, command):

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

396

397 # 管道检查

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

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

400

401 # 重定向检查

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

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

404

405 return ""

406

407 @staticmethod

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

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

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

411 if len(encoded) <= max_bytes:

412 return text

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

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

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

152 statements