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 等) 直接拒绝。
133 lines
4.7 KiB
Python
133 lines
4.7 KiB
Python
import unittest
|
|
|
|
from backend.agent_tools import (
|
|
classify_agent_command,
|
|
classify_kubectl_command,
|
|
classify_shell_command,
|
|
parse_kubectl_command,
|
|
parse_shell_command,
|
|
)
|
|
|
|
|
|
class KubectlCommandPolicyTest(unittest.TestCase):
|
|
def test_allows_read_only_kubectl_command_for_autonomous_execution(self):
|
|
decision = classify_kubectl_command("kubectl get pods -n production -o wide")
|
|
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_requires_approval_for_mutating_kubectl_command(self):
|
|
decision = classify_kubectl_command("kubectl rollout restart deployment/api -n production")
|
|
|
|
self.assertEqual(decision.mode, "write")
|
|
self.assertTrue(decision.requires_approval)
|
|
|
|
def test_rejects_shell_operators_and_non_kubectl_commands(self):
|
|
for command in (
|
|
"kubectl get pods; rm -rf /",
|
|
"kubectl get pods | sh",
|
|
"bash -c kubectl get pods",
|
|
"sudo kubectl get pods",
|
|
):
|
|
with self.subTest(command=command):
|
|
with self.assertRaises(ValueError):
|
|
parse_kubectl_command(command)
|
|
|
|
def test_rejects_exec_even_when_command_looks_read_only(self):
|
|
decision = classify_kubectl_command("kubectl exec api-0 -- cat /etc/hosts")
|
|
|
|
self.assertEqual(decision.mode, "write")
|
|
self.assertTrue(decision.requires_approval)
|
|
|
|
|
|
class ShellCommandPolicyTest(unittest.TestCase):
|
|
def test_allows_read_only_systemctl_status_for_autonomous_execution(self):
|
|
decision = classify_shell_command("systemctl status kubelet")
|
|
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_allows_journalctl_with_flags(self):
|
|
decision = classify_shell_command("journalctl -u kubelet -n 100 --no-pager")
|
|
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_allows_crictl_ps(self):
|
|
decision = classify_shell_command("crictl ps -a")
|
|
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_allows_df_and_free(self):
|
|
for command in ("df -h", "free -m", "uptime"):
|
|
with self.subTest(command=command):
|
|
decision = classify_shell_command(command)
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_allows_ip_addr(self):
|
|
decision = classify_shell_command("ip addr show")
|
|
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_requires_approval_for_systemctl_restart(self):
|
|
decision = classify_shell_command("systemctl restart kubelet")
|
|
|
|
self.assertEqual(decision.mode, "write")
|
|
self.assertTrue(decision.requires_approval)
|
|
|
|
def test_rejects_unknown_binary(self):
|
|
with self.assertRaises(ValueError):
|
|
classify_shell_command("rm -rf /")
|
|
|
|
def test_rejects_shell_operators_in_shell_command(self):
|
|
for command in (
|
|
"systemctl status kubelet; rm -rf /",
|
|
"df -h | grep vda",
|
|
"journalctl -u kubelet > /tmp/log",
|
|
):
|
|
with self.subTest(command=command):
|
|
with self.assertRaises(ValueError):
|
|
parse_shell_command(command)
|
|
|
|
def test_cat_only_allows_safe_paths(self):
|
|
decision = classify_shell_command("cat /etc/os-release")
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_cat_rejects_unsafe_path(self):
|
|
with self.assertRaises(ValueError):
|
|
classify_shell_command("cat /etc/shadow")
|
|
|
|
|
|
class AgentCommandDispatchTest(unittest.TestCase):
|
|
def test_dispatches_kubectl_command_to_kubectl_classifier(self):
|
|
decision = classify_agent_command("kubectl get pods -A")
|
|
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertEqual(decision.verb, "get")
|
|
|
|
def test_dispatches_shell_command_to_shell_classifier(self):
|
|
decision = classify_agent_command("systemctl status kubelet")
|
|
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertEqual(decision.verb, "systemctl")
|
|
|
|
def test_dispatches_write_shell_command_to_shell_classifier(self):
|
|
decision = classify_agent_command("systemctl restart kubelet")
|
|
|
|
self.assertEqual(decision.mode, "write")
|
|
self.assertTrue(decision.requires_approval)
|
|
|
|
def test_dispatches_write_kubectl_command(self):
|
|
decision = classify_agent_command("kubectl delete pod api-0 -n prod")
|
|
|
|
self.assertEqual(decision.mode, "write")
|
|
self.assertTrue(decision.requires_approval)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|