diff --git a/backend/agent_tools.py b/backend/agent_tools.py index 917e782..6d1b54b 100644 --- a/backend/agent_tools.py +++ b/backend/agent_tools.py @@ -31,12 +31,11 @@ GLOBAL_FLAGS_WITH_VALUE = { "--namespace", "-n", "--profile", "--request-timeout", "--server", "-s", "--tls-server-name", "--token", "--user", } -FORBIDDEN_TOKENS = {";", "||", "&&", "<", "<<", "`"} -# 管道 |、重定向 >、>>、$ 变量引用 允许在诊断中使用 -# 但含 | 的命令中若管道后出现 sh/bash/dd/format/mkfs/fdisk/parted 则直接拒绝（武器级危险） -# rm/mv/chmod/chown — 现已在 SHELL_WRITE_COMMANDS 中，管道后会触发审批而非直接拒绝 -_PIPE_DANGEROUS_TARGETS = {"sh", "bash", "dash", "zsh", - "dd", "format", "mkfs", "fdisk", "parted"} +FORBIDDEN_TOKENS = set() # 全部 shell 操作符均允许；安全靠白名单+危险命令检查 +# 永远拒绝的武器级危险命令（无论是否在白名单中） +_ALWAYS_DENY_CMDS = {"sh", "bash", "dash", "zsh", + "dd", "format", "mkfs", "fdisk", "parted"} + @dataclass(frozen=True) @@ -49,14 +48,10 @@ class CommandDecision: def _has_shell_operators(command: str) -> bool: - """检测命令是否包含管道 | 或重定向 >。""" - # 简单检测：提取所有 token 中的操作符 - # 注意排除出现在引号内的 | 和 > + """检测命令是否包含 shell 操作符（管道、重定向、分号、逻辑连接符等）。""" 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: @@ -65,33 +60,22 @@ def _has_shell_operators(command: str) -> bool: 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 # > 或 >> 允许 + if ch in ('|', '>', ';', '&', '<'): + # 单纯 & 不算（后台任务标记） + if ch == '&': + if i + 1 < len(stripped) and stripped[i + 1] in ('&',): + return True # && + continue + return True return False -def _pipe_to_dangerous_target(command: str) -> bool: - """检查管道后面的命令是否是危险命令。""" - if '|' not in command: - return False - # 按管道分割，取最后一段 - # 简单解析：用 | 分割，忽略引号内的 - parts = [] +def _split_into_commands(command: str) -> list[str]: + """按 shell 连接符（; && || | > >> < << 等）分割命令块。 + + 返回每个可执行段，忽略引号内的操作符。用于白名单检查。 + """ + segments = [] buf = [] in_single = in_double = False for ch in command: @@ -99,30 +83,28 @@ def _pipe_to_dangerous_target(command: str) -> bool: 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:]: + if not in_single and not in_double and ch in ('|', ';', '&', '>', '<'): + # 包含重定向符号时，前面的部分也是可执行段 + if buf: + segments.append(''.join(buf).strip()) + buf = [] + if ch == '&': + continue # & 符号自身不产生段 + continue + buf.append(ch) + if buf: + segments.append(''.join(buf).strip()) + return [s for s in segments if s] + + +def _command_has_dangerous_cmds(command: str) -> bool: + """检查命令中任何可执行段是否包含武器级危险命令。""" + for segment in _split_into_commands(command): try: - part_args = shlex.split(part.strip()) + seg_args = shlex.split(segment) except ValueError: continue - if part_args and part_args[0] in _PIPE_DANGEROUS_TARGETS: + if seg_args and seg_args[0] in _ALWAYS_DENY_CMDS: return True return False @@ -130,18 +112,15 @@ def _pipe_to_dangerous_target(command: str) -> bool: def _parse_command(command: str) -> list[str]: """解析命令，检查合法性。 - 允许 kubectl 和诊断命令中的管道 (|)、重定向 (>)、$ 变量引用。 - 不允许的: ; || && < << 后台命令 ` 以及多行。 - 管道后不允许出现 sh/bash/rm/mv/chown/chmod/dd 等危险命令。 + 允许所有 shell 操作符（| > ; && || < 等），安全靠白名单检查。 + 武器级危险命令 sh/bash/dd/mkfs 等在 classify 阶段处理。 """ if not isinstance(command, str) or not command.strip(): raise ValueError("命令不能为空") 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 等），如需要请手动执行") + if _command_has_dangerous_cmds(command): + raise ValueError("命令包含不允许的危险命令（sh/bash/dd/mkfs 等），请手动执行") try: return shlex.split(command) except ValueError as exc: @@ -229,6 +208,8 @@ SHELL_READ_ONLY_COMMANDS: dict[str, set[str] | None] = { "curl": None, "wget": None, "traceroute": None, "tracepath": None, "mtr": None, "nc": None, "nmap": None, "telnet": None, "tcpdump": None, "ethtool": None, "lspci": None, "lsusb": None, "dmidecode": None, "sysctl": None, "udevadm": None, + "echo": None, "printf": None, "cd": None, "pwd": None, "export": None, + "test": None, "[": None, } # etcdctl 由 classify_shell_command 中的专用块处理 (两段式命令) @@ -313,42 +294,37 @@ def classify_shell_command(command: str) -> CommandDecision: 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() + # 含 shell 操作符：先检查武器级危险命令 + if _command_has_dangerous_cmds(command): + raise ValueError("命令包含不允许的危险命令（sh/bash/dd/mkfs 等），请手动执行") + # 取第一个可执行段的 binary 做白名单检查 + segments = _split_into_commands(command) + first_seg = segments[0] if segments else command try: - before_parts = shlex.split(before_pipe) + first_parts = shlex.split(first_seg) except ValueError: - before_parts = [before_pipe] - binary = before_parts[0] + first_parts = [first_seg] + binary = first_parts[0] if binary not in SHELL_READ_ONLY_COMMANDS and binary not in SHELL_WRITE_COMMANDS: raise ValueError(f"不允许执行该命令: {binary}") - # cat 的安全路径检查（管道场景） + # cat 的安全路径检查 if binary == "cat": - for arg in before_parts[1:]: + for arg in first_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:]) + verb = _first_non_flag(first_parts[1:]) if verb is not None and verb in allowed: return CommandDecision("write", True, binary, [command], needs_shell=True) - # binary 在写表但子命令不匹配写表 — 走只读白名单检查 + # 写表子命令不匹配 — 看是否在只读表中 if binary in SHELL_READ_ONLY_COMMANDS: ro_allowed = SHELL_READ_ONLY_COMMANDS[binary] if ro_allowed is None or (verb is not None and verb in ro_allowed): diff --git a/backend/ai_agent.py b/backend/ai_agent.py index 3551a11..6d70659 100644 --- a/backend/ai_agent.py +++ b/backend/ai_agent.py @@ -40,7 +40,7 @@ AGENT_SYSTEM_PROMPT = """你是 Kubernetes 自主诊断 Agent。你的首要职 3. 每次只提出一条命令，服务端会执行并把真实输出返回给你，然后你继续下一轮判断。 安全规则： -1. 每次只返回一条命令，但可以使用管道 (|) 连接命令进行精确过滤、使用重定向 (>) 保存输出到文件。禁止使用 ; || && 连接多条不相关的命令。 +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 f85e4f2..5aea57d 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -23,10 +23,7 @@ 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 等危险命令仍被拒绝 + # 武器级危险命令（sh/bash）仍被拒绝 with self.assertRaises(ValueError): parse_kubectl_command("kubectl get pods | sh") # 非 kubectl 命令仍被拒绝 @@ -93,18 +90,26 @@ class ShellCommandPolicyTest(unittest.TestCase): classify_shell_command("hacktool -x") def test_rejects_shell_operators_in_shell_command(self): - # 分号 ; 仍被拒绝 - with self.assertRaises(ValueError): - parse_shell_command("systemctl status kubelet; rm -rf /") - # 管道 | 和重定向 > 现在允许通过分类器（server 端会用 shell=True 执行） + # 所有 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")