Your Name
00d95c4c5e
现在 AI Agent 支持在单条命令中使用: - 管道 (|): crictl ps -a | grep nginx | head -5 - 重定向 (> / >>): journalctl -u kubelet > /tmp/log - 分号 (;): systemctl status kubelet; df -h - 逻辑连接符 (&& / ||): df -h && free -m 安全策略: - FORBIDDEN_TOKENS 清空，所有 shell 操作符均放行 - 新增 _split_into_commands() 按操作符分割命令段 - 新增 _command_has_dangerous_cmds() 扫描所有段中的武器级命令 - 永远拒绝: sh/bash/dash/zsh/dd/format/.../parted 等 - 检查第一条命令的白名单，cat 不安全路径检查 - 只读白名单新增 echo printf cd pwd export test [
212 lines
8.5 KiB
Python
212 lines
8.5 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):
|
|
# 武器级危险命令（sh/bash）仍被拒绝
|
|
with self.assertRaises(ValueError):
|
|
parse_kubectl_command("kubectl get pods | sh")
|
|
# 非 kubectl 命令仍被拒绝
|
|
for command in (
|
|
"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):
|
|
# rm 现在在写白名单中（需审批），不会直接拒绝
|
|
decision = classify_shell_command("rm -rf /tmp/cache")
|
|
self.assertEqual(decision.mode, "write")
|
|
self.assertTrue(decision.requires_approval)
|
|
# 完全不在任何白名单中的命令才拒绝
|
|
with self.assertRaises(ValueError):
|
|
classify_shell_command("hacktool -x")
|
|
|
|
def test_rejects_shell_operators_in_shell_command(self):
|
|
# 所有 shell 操作符（; | > && ||）现在都允许通过分类器
|
|
# 分号 ;
|
|
decision = classify_shell_command("systemctl status kubelet; df -h")
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertTrue(decision.needs_shell)
|
|
# 管道 |
|
|
decision = classify_shell_command("df -h | grep vda")
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
self.assertTrue(decision.needs_shell)
|
|
# 重定向 >
|
|
decision = classify_shell_command("journalctl -u kubelet > /tmp/log")
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
self.assertTrue(decision.needs_shell)
|
|
# 武器级 sh/bash 仍拒绝
|
|
with self.assertRaises(ValueError):
|
|
classify_shell_command("crictl ps -a | bash")
|
|
with self.assertRaises(ValueError):
|
|
classify_shell_command("crictl ps -a && sh -c 'echo x'")
|
|
|
|
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 EtcdCommandPolicyTest(unittest.TestCase):
|
|
def test_etcdctl_endpoint_status_is_read_only(self):
|
|
decision = classify_shell_command("etcdctl endpoint status")
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_etcdctl_with_global_flags_still_classifies(self):
|
|
decision = classify_shell_command(
|
|
"etcdctl --endpoints=https://127.0.0.1:2379 --cacert=/etc/ssl/ca.pem endpoint health"
|
|
)
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_etcdctl_member_list_is_read_only(self):
|
|
decision = classify_shell_command("etcdctl member list")
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_etcdctl_get_and_version_are_read_only(self):
|
|
for command in ("etcdctl get /kubernetes", "etcdctl version", "etcdctl alarm list"):
|
|
with self.subTest(command=command):
|
|
decision = classify_shell_command(command)
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_etcdctl_alarm_list_is_read_only(self):
|
|
decision = classify_shell_command("etcdctl alarm list")
|
|
self.assertEqual(decision.mode, "read")
|
|
self.assertFalse(decision.requires_approval)
|
|
|
|
def test_etcdctl_put_requires_approval(self):
|
|
decision = classify_shell_command("etcdctl put key value")
|
|
self.assertEqual(decision.mode, "write")
|
|
self.assertTrue(decision.requires_approval)
|
|
|
|
def test_etcdctl_member_remove_requires_approval(self):
|
|
decision = classify_shell_command("etcdctl member remove abc123")
|
|
self.assertEqual(decision.mode, "write")
|
|
self.assertTrue(decision.requires_approval)
|
|
|
|
def test_etcdctl_snapshot_save_requires_approval(self):
|
|
decision = classify_shell_command("etcdctl snapshot save /tmp/etcd.db")
|
|
self.assertEqual(decision.mode, "write")
|
|
self.assertTrue(decision.requires_approval)
|
|
|
|
def test_etcdctl_compact_requires_approval(self):
|
|
decision = classify_shell_command("etcdctl compact 12345")
|
|
self.assertEqual(decision.mode, "write")
|
|
self.assertTrue(decision.requires_approval)
|
|
|
|
def test_etcdctl_auth_enable_requires_approval(self):
|
|
decision = classify_shell_command("etcdctl auth enable")
|
|
self.assertEqual(decision.mode, "write")
|
|
self.assertTrue(decision.requires_approval)
|
|
|
|
def test_etcdctl_rejects_unknown_subcommand(self):
|
|
with self.assertRaises(ValueError):
|
|
classify_shell_command("etcdctl bogus command")
|
|
|
|
|
|
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()
|