cnbugs
32bea0d4ea
扩展 Agent 命令工具，除 kubectl 外新增节点宿主机系统诊断命令支持：
- agent_tools.py: 新增 shell 命令白名单分类器 (classify_shell_command)
- 只读白名单: systemctl status/is-active、journalctl、crictl ps/logs、
docker ps/logs、df、free、ss、ip addr、ps、mount、lsmod、dmesg、
ping、nslookup、cat (限 /proc /sys /etc/kubernetes 等安全路径) 等 30+ 命令
- 写白名单: systemctl restart/stop/start/reload 等需人工审批
- 统一分发器 classify_agent_command 按命令前缀路由
- 统一执行入口 execute_command 走 subprocess_exec (无 shell 注入风险)
- ai_agent.py: 更新 AGENT_SYSTEM_PROMPT 告知 AI 两类可用工具及安全规则
- 三个调用点 (初始探测/分类/执行) 切换到统一接口
- main.py: 审批执行处改用 execute_command
- tests: 新增 19 个用例覆盖 shell 命令分类、混合执行、白名单拒绝
- 全部 44 个测试通过
- 前端/README: 文案与文档补充 shell 命令支持说明
安全策略不变: 禁止 shell 操作符/管道/重定向/多命令拼接，
非白名单命令 (rm/vi/apt 等) 直接拒绝。
98 lines
4.5 KiB
Python
98 lines
4.5 KiB
Python
import asyncio
|
|
import json
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import backend.main as main
|
|
|
|
|
|
def parse_sse(event):
|
|
return json.loads(event[6:].strip())
|
|
|
|
|
|
class AgentApiTest(unittest.TestCase):
|
|
def test_agent_stream_forwards_agent_events(self):
|
|
async def fake_agent(report, question="", max_steps=6, execution_mode="ai"):
|
|
self.assertEqual(execution_mode, "manual")
|
|
yield {"type": "thinking", "step": 1, "message": "规划"}
|
|
yield {"type": "final", "content": "完成"}
|
|
|
|
async def run():
|
|
main._last_diagnosis = {"score": 60}
|
|
main._last_report = "report"
|
|
with patch.object(main, "run_diagnostic_agent", fake_agent):
|
|
response = await main.agent_diagnose_stream(question="检查 Pod", execution_mode="manual")
|
|
return [parse_sse(event) async for event in response.body_iterator]
|
|
|
|
events = asyncio.run(run())
|
|
self.assertEqual(events[0]["type"], "start")
|
|
self.assertEqual(events[-1]["type"], "done")
|
|
self.assertEqual(events[-2]["type"], "final")
|
|
|
|
def test_config_view_is_read_only_and_does_not_require_approval(self):
|
|
from backend.agent_tools import classify_kubectl_command
|
|
|
|
decision = classify_kubectl_command("kubectl config view")
|
|
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_execute_approved_command_continues_agent_with_command_result(self):
|
|
approval = {
|
|
"command": "kubectl rollout restart deployment/api -n prod",
|
|
"reason": "重启工作负载",
|
|
"run_id": "",
|
|
"report": "diagnosis report",
|
|
"question": "修复 API",
|
|
"execution_mode": "ai",
|
|
"messages": [{"role": "system", "content": "context"}],
|
|
"next_step": 2,
|
|
"max_steps": 30,
|
|
"command_counts": {"kubectl rollout restart deployment/api -n prod": 1},
|
|
}
|
|
|
|
async def fake_execute(command):
|
|
return {"command": command, "mode": "write", "returncode": 0, "stdout": "restarted", "stderr": ""}
|
|
|
|
async def fake_resume(pending, result):
|
|
self.assertEqual(result["stdout"], "restarted")
|
|
yield {"type": "thinking", "step": 2, "message": "继续验证"}
|
|
yield {"type": "command", "command": "kubectl get pods -n prod", "mode": "read", "reason": "验证恢复"}
|
|
yield {"type": "result", "command": "kubectl get pods -n prod", "mode": "read", "returncode": 0, "stdout": "Running", "stderr": ""}
|
|
yield {"type": "final", "content": "故障已恢复", "model": "test"}
|
|
|
|
async def run():
|
|
with patch.object(main, "take_pending_approval", return_value=approval), patch.object(main, "execute_command", fake_execute), patch.object(main, "resume_diagnostic_agent", fake_resume):
|
|
response = await main.execute_approved_agent_command(main.AgentApprovalRequest(approval_id="abc", approved=True))
|
|
return [parse_sse(event) async for event in response.body_iterator]
|
|
|
|
events = asyncio.run(run())
|
|
self.assertEqual(events[0]["type"], "approved_result")
|
|
self.assertEqual(events[1]["type"], "thinking")
|
|
self.assertEqual(events[-2]["type"], "final")
|
|
self.assertEqual(events[-1]["type"], "done")
|
|
|
|
def test_execute_approved_command_once(self):
|
|
approval = {"command": "kubectl scale deployment/api --replicas=2", "reason": "恢复副本"}
|
|
|
|
async def fake_execute(command):
|
|
return {"command": command, "mode": "write", "returncode": 0, "stdout": "scaled", "stderr": ""}
|
|
|
|
async def fake_resume(pending, result):
|
|
yield {"type": "final", "content": "验证完成", "model": "test"}
|
|
|
|
async def run():
|
|
with patch.object(main, "take_pending_approval", return_value=approval), patch.object(main, "execute_command", fake_execute), patch.object(main, "resume_diagnostic_agent", fake_resume):
|
|
response = await main.execute_approved_agent_command(main.AgentApprovalRequest(approval_id="abc", approved=True))
|
|
return [parse_sse(event) async for event in response.body_iterator]
|
|
|
|
events = asyncio.run(run())
|
|
self.assertEqual(events[0]["type"], "approved_result")
|
|
self.assertEqual(events[0]["returncode"], 0)
|
|
self.assertEqual(events[0]["stdout"], "scaled")
|
|
self.assertEqual(events[-1]["type"], "done")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|