diff --git a/backend/agent_tools.py b/backend/agent_tools.py index e3259b0..bb0fe41 100644 --- a/backend/agent_tools.py +++ b/backend/agent_tools.py @@ -31,7 +31,11 @@ GLOBAL_FLAGS_WITH_VALUE = { "--namespace", "-n", "--profile", "--request-timeout", "--server", "-s", "--tls-server-name", "--token", "--user", } -FORBIDDEN_TOKENS = {";", "|", "||", "&&", ">", ">>", "<", "<<", "`", "$"} +FORBIDDEN_TOKENS = {";", "||", "&&", "<", "<<", "`"} +# 管道 |、重定向 >、>>、$ 变量引用 允许在诊断中使用 +# 但含 | 的命令中若管道后出现 sh/bash/rm/mv/chmod/chown/dd/format/mkfs 则拒绝 +_PIPE_DANGEROUS_TARGETS = {"sh", "bash", "dash", "zsh", "rm", "mv", "chmod", "chown", + "dd", "format", "mkfs", "fdisk", "parted", "wget", "curl"} @dataclass(frozen=True) @@ -40,14 +44,103 @@ class CommandDecision: requires_approval: bool verb: str # kubectl 子命令 或 shell 命令名 args: list[str] + needs_shell: bool = False # True 时需用 shell=True 执行 + + +def _has_shell_operators(command: str) -> bool: + """检测命令是否包含管道 | 或重定向 >。""" + # 简单检测:提取所有 token 中的操作符 + # 注意排除出现在引号内的 | 和 > + stripped = command.strip() + if not stripped: + return False + # 先用 shlex 分割看是否含有 shell 操作符作为独立 token + # 但更可靠的是直接扫描字符(考虑引号保护) + in_single = in_double = False + for i, ch in enumerate(stripped): + if ch == "'" and not in_double: + in_single = not in_single + elif ch == '"' and not in_single: + in_double = not in_double + if in_single or in_double: + continue + if ch in ('|', '>', ';', '&'): + # ; 和 && 仍然不允许 + if ch in (';', '&'): + # 单独 & 不算(后台),&& 才不允许 + if ch == ';': + return True + # 检查是否是 && + if ch == '&' and i + 1 < len(stripped) and stripped[i + 1] == '&': + return True + if ch == '|': + # | 允许;但检查是否是 || + if i + 1 < len(stripped) and stripped[i + 1] == '|': + # || 不允许 + return False + return True # 单 | 是管道,允许 + if ch == '>': + return True # > 或 >> 允许 + return False + + +def _pipe_to_dangerous_target(command: str) -> bool: + """检查管道后面的命令是否是危险命令。""" + if '|' not in command: + return False + # 按管道分割,取最后一段 + # 简单解析:用 | 分割,忽略引号内的 + parts = [] + buf = [] + in_single = in_double = False + for ch in command: + if ch == "'" and not in_double: + in_single = not in_single + elif ch == '"' and not in_single: + in_double = not in_double + if ch == '|' and not in_single and not in_double: + parts.append(''.join(buf)) + buf = [] + else: + buf.append(ch) + parts.append(''.join(buf)) + if len(parts) < 2: + return False + # 检查管道最后一条命令的危险词 + last_part = parts[-1].strip() + # shlex 取第一个词 + try: + last_args = shlex.split(last_part) + except ValueError: + return False + if last_args and last_args[0] in _PIPE_DANGEROUS_TARGETS: + return True + # 也检查中间管道(如 docker exec | sh) + for part in parts[1:]: + try: + part_args = shlex.split(part.strip()) + except ValueError: + continue + if part_args and part_args[0] in _PIPE_DANGEROUS_TARGETS: + return True + return False def _parse_command(command: str) -> list[str]: - """解析单条命令,禁止 shell 操作符、重定向和命令组合。""" + """解析命令,检查合法性。 + + 允许 kubectl 和诊断命令中的管道 (|)、重定向 (>)、$ 变量引用。 + 不允许的: ; || && < << 后台命令 ` 以及多行。 + 管道后不允许出现 sh/bash/rm/mv/chown/chmod/dd 等危险命令。 + """ if not isinstance(command, str) or not command.strip(): raise ValueError("命令不能为空") - if any(token in command for token in FORBIDDEN_TOKENS) or "\n" in command or "\r" in command: - raise ValueError("不允许 shell 操作符、重定向或多条命令") + if "\n" in command or "\r" in command: + raise ValueError("不允许多行命令") + if any(token in command for token in FORBIDDEN_TOKENS): + raise ValueError("不允许 shell 命令连接符(; || &&)或输入重定向(< <<)") + if _pipe_to_dangerous_target(command): + raise ValueError("管道后不允许执行危险命令(sh/rm/chmod 等),如需要请手动执行") try: return shlex.split(command) except ValueError as exc: @@ -203,7 +296,51 @@ def classify_shell_command(command: str) -> CommandDecision: systemctl 同时出现在只读和写白名单中:restart/stop 等走写表(需审批), status/is-active 等走只读表(自动执行)。因此先检查写表,再检查只读表。 + + 含管道 (|) 或重定向 (>) 的命令:只检查管道前的 binary 是否在白名单内, + 并设置 needs_shell=True 以便用 shell 方式执行。 """ + needs_shell = _has_shell_operators(command) + + if needs_shell: + # 含 shell 操作符:排除分号(分号仍被禁止) + if ';' in command: + raise ValueError("不允许使用分号(;)连接多条命令") + # 含管道/重定向:先检查管道后是否有危险命令 + if _pipe_to_dangerous_target(command): + raise ValueError("管道后不允许执行危险命令(sh/rm/chmod 等),如需要请手动执行") + # 取管道前的 binary 做白名单检查 + pipe_idx = command.index("|") if "|" in command else len(command) + before_pipe = command[:pipe_idx].strip() + try: + before_parts = shlex.split(before_pipe) + except ValueError: + before_parts = [before_pipe] + binary = before_parts[0] + + if binary not in SHELL_READ_ONLY_COMMANDS and binary not in SHELL_WRITE_COMMANDS: + raise ValueError(f"不允许执行该命令: {binary}") + # cat 的安全路径检查(管道场景) + if binary == "cat": + for arg in before_parts[1:]: + if not arg.startswith("-") and not _is_safe_cat_path(arg): + raise ValueError( + "cat 仅允许读取 /proc, /sys, /etc/kubernetes, /etc/cni, " + "/etc/systemd/system, /etc/os-release, /etc/hosts, /etc/resolv.conf" + ) + # 对于含管道的命令,只检查 binary 白名单,不做子命令级别的检查 + # 同时拒绝写操作(如 systemctl restart kubelet | grep ... 也不允许自动执行) + if binary in SHELL_WRITE_COMMANDS: + allowed = SHELL_WRITE_COMMANDS[binary] + if allowed is None: + return CommandDecision("write", True, binary, [command], needs_shell=True) + # 取管道前的子命令 + verb = _first_non_flag(before_parts[1:]) + if verb is not None and verb in allowed: + return CommandDecision("write", True, binary, [command], needs_shell=True) + + return CommandDecision("read", False, binary, [command], needs_shell=True) + parts = parse_shell_command(command) binary = parts[0] verb = _first_non_flag(parts[1:]) @@ -262,7 +399,10 @@ def classify_agent_command(command: str) -> CommandDecision: # ── 命令执行 ────────────────────────────────────────────── async def execute_command(command: str, timeout: int = 30) -> dict: - """执行单条 Agent 命令 (kubectl 或白名单系统诊断命令),无 shell。""" + """执行 Agent 命令 (kubectl 或白名单系统诊断命令)。 + + 含管道/重定向的命令通过 shell=True 执行,其他用 subprocess_exec 直接执行。 + """ decision = classify_agent_command(command) parts = _parse_command(command) if parts[0] == "kubectl": @@ -272,12 +412,19 @@ async def execute_command(command: str, timeout: int = 30) -> dict: executable = parts[0] args = parts[1:] - proc = await asyncio.create_subprocess_exec( - executable, - *args, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) + if decision.needs_shell: + proc = await asyncio.create_subprocess_shell( + command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + else: + proc = await asyncio.create_subprocess_exec( + executable, + *args, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) try: stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout) return { diff --git a/backend/ai_agent.py b/backend/ai_agent.py index 83eaa56..3551a11 100644 --- a/backend/ai_agent.py +++ b/backend/ai_agent.py @@ -40,7 +40,7 @@ AGENT_SYSTEM_PROMPT = """你是 Kubernetes 自主诊断 Agent。你的首要职 3. 每次只提出一条命令,服务端会执行并把真实输出返回给你,然后你继续下一轮判断。 安全规则: -1. 每次只返回一条命令,不得使用 shell、管道、重定向或命令连接符 (; | && > < 等)。 +1. 每次只返回一条命令,但可以使用管道 (|) 连接命令进行精确过滤、使用重定向 (>) 保存输出到文件。禁止使用 ; || && 连接多条不相关的命令。 2. kubectl 的 get、describe、logs、top 等只读命令可以自动执行。 3. 系统诊断只读命令 (systemctl status、journalctl、crictl ps、df、free、ss、ip addr 等) 可以自动执行。 4. apply、delete、patch、scale、rollout、exec 等修改或交互命令必须暂停并请求人工批准。 diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 02d6b3a..ef254ea 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -23,9 +23,14 @@ class KubectlCommandPolicyTest(unittest.TestCase): self.assertTrue(decision.requires_approval) def test_rejects_shell_operators_and_non_kubectl_commands(self): + # 分号 ; 仍被拒绝 + with self.assertRaises(ValueError): + parse_kubectl_command("kubectl get pods; rm -rf /") + # 管道后的 sh/bash 等危险命令仍被拒绝 + with self.assertRaises(ValueError): + parse_kubectl_command("kubectl get pods | sh") + # 非 kubectl 命令仍被拒绝 for command in ( - "kubectl get pods; rm -rf /", - "kubectl get pods | sh", "bash -c kubectl get pods", "sudo kubectl get pods", ): @@ -83,14 +88,18 @@ class ShellCommandPolicyTest(unittest.TestCase): 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) + # 分号 ; 仍被拒绝 + with self.assertRaises(ValueError): + parse_shell_command("systemctl status kubelet; rm -rf /") + # 管道 | 和重定向 > 现在允许通过分类器(server 端会用 shell=True 执行) + 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) def test_cat_only_allows_safe_paths(self): decision = classify_shell_command("cat /etc/os-release")