k8smanager-cli/tests/test_agent_tools.py

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):
        # cat 完全放行，任意路径均可
        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):
        # cat 不再限制路径
        decision = classify_shell_command("cat /etc/shadow")
        self.assertEqual(decision.mode, "read")
        self.assertFalse(decision.requires_approval)


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()