cnbugs
110 lines
3.8 KiB
Python
110 lines
3.8 KiB
Python
"""AI Agent 可用的 kubectl 命令工具和安全策略。"""
|
|
import asyncio
|
|
import shlex
|
|
from dataclasses import dataclass
|
|
|
|
from backend.diagnosis import KUBECTL
|
|
|
|
|
|
READ_ONLY_VERBS = {
|
|
"api-resources", "api-versions", "auth", "cluster-info", "describe",
|
|
"diff", "explain", "get", "logs", "top", "version", "wait",
|
|
}
|
|
WRITE_OR_INTERACTIVE_VERBS = {
|
|
"annotate", "apply", "attach", "autoscale", "cordon", "cp", "create",
|
|
"debug", "delete", "drain", "edit", "exec", "expose", "label", "patch",
|
|
"port-forward", "replace", "rollout", "run", "scale", "set", "taint",
|
|
"uncordon",
|
|
}
|
|
GLOBAL_FLAGS_WITH_VALUE = {
|
|
"--as", "--as-group", "--as-uid", "--cache-dir", "--certificate-authority",
|
|
"--client-certificate", "--client-key", "--cluster", "--context", "--kubeconfig",
|
|
"--namespace", "-n", "--profile", "--request-timeout", "--server", "-s",
|
|
"--tls-server-name", "--token", "--user",
|
|
}
|
|
FORBIDDEN_TOKENS = {";", "|", "||", "&&", ">", ">>", "<", "<<", "`", "$"}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CommandDecision:
|
|
mode: str
|
|
requires_approval: bool
|
|
verb: str
|
|
args: list[str]
|
|
|
|
|
|
def parse_kubectl_command(command: str) -> list[str]:
|
|
"""仅解析单条 kubectl 命令，不允许 shell、重定向或命令组合。"""
|
|
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 操作符、重定向或多条命令")
|
|
try:
|
|
parts = shlex.split(command)
|
|
except ValueError as exc:
|
|
raise ValueError(f"命令格式错误: {exc}") from exc
|
|
if not parts or parts[0] != "kubectl":
|
|
raise ValueError("只允许执行 kubectl 命令")
|
|
if len(parts) < 2:
|
|
raise ValueError("kubectl 命令缺少操作类型")
|
|
return parts[1:]
|
|
|
|
|
|
def _find_verb(args: list[str]) -> str:
|
|
index = 0
|
|
while index < len(args):
|
|
arg = args[index]
|
|
if arg == "--":
|
|
break
|
|
if arg in GLOBAL_FLAGS_WITH_VALUE:
|
|
index += 2
|
|
continue
|
|
if any(arg.startswith(flag + "=") for flag in GLOBAL_FLAGS_WITH_VALUE if flag.startswith("--")):
|
|
index += 1
|
|
continue
|
|
if arg.startswith("-"):
|
|
index += 1
|
|
continue
|
|
return arg.lower()
|
|
raise ValueError("无法识别 kubectl 操作类型")
|
|
|
|
|
|
def classify_kubectl_command(command: str) -> CommandDecision:
|
|
args = parse_kubectl_command(command)
|
|
verb = _find_verb(args)
|
|
if verb in READ_ONLY_VERBS:
|
|
return CommandDecision("read", False, verb, args)
|
|
if verb in WRITE_OR_INTERACTIVE_VERBS:
|
|
return CommandDecision("write", True, verb, args)
|
|
return CommandDecision("write", True, verb, args)
|
|
|
|
|
|
async def execute_kubectl(command: str, timeout: int = 30) -> dict:
|
|
"""无 shell 执行 kubectl，并限制输出大小。"""
|
|
decision = classify_kubectl_command(command)
|
|
proc = await asyncio.create_subprocess_exec(
|
|
KUBECTL,
|
|
*decision.args,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
try:
|
|
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
|
|
return {
|
|
"command": command,
|
|
"mode": decision.mode,
|
|
"returncode": proc.returncode,
|
|
"stdout": stdout.decode(errors="replace")[:20000].strip(),
|
|
"stderr": stderr.decode(errors="replace")[:10000].strip(),
|
|
}
|
|
except asyncio.TimeoutError:
|
|
proc.kill()
|
|
await proc.communicate()
|
|
return {
|
|
"command": command,
|
|
"mode": decision.mode,
|
|
"returncode": -1,
|
|
"stdout": "",
|
|
"stderr": f"命令执行超时 ({timeout}s)",
|
|
}
|