langchain-agentx-python 0.1__py3-none-any.whl

langchain_agentx/tools/bash/ast_security.py

@@ -0,0 +1,571 @@
+"""
+tools/bash/ast_security.py — BashRuntimeTool v2 AST 分析
+
+职责：
+    使用 tree-sitter 对 bash 命令做结构化分析，提供：
+    - 顶层子命令拆分（pipeline / list）
+    - 命令替换 / 进程替换 / subshell / command group 检测
+    - 包装命令归一化（env / command / builtin / nohup）
+    - 嵌套 shell 执行检测（bash -c / sh -c / zsh -c）
+    - 输出重定向检测（辅助只读/写入判定）
+
+对照 CC：
+    - bashSecurity.ts                    → 结构级危险模式识别
+    - bashCommandHelpers.ts              → 子命令拆分与分段授权
+    - ParsedCommand / treeSitterAnalysis → 结构化命令理解
+
+实现约束：
+    - 采用 OOP 风格，避免散落函数主导的实现
+    - 保持 fail-open：AST 失败时回退到轻量启发式，而不是阻塞工具
+"""
+
+from __future__ import annotations
+
+import re
+import shlex
+from dataclasses import dataclass
+from functools import lru_cache
+
+from tree_sitter import Node
+from tree_sitter_languages import get_parser
+
+
+@dataclass(frozen=True)
+class BashSimpleCommand:
+    """单个 simple command 的结构摘要。"""
+
+    text: str
+    argv: tuple[str, ...]
+    base_command: str | None
+    normalized_base_command: str | None
+
+
+@dataclass(frozen=True)
+class BashAstAnalysis:
+    """
+    Bash AST 分析结果。
+
+    `permission_segments` 是 v2 最关键字段：用于对顶层 pipeline / list
+    做逐段授权，对齐 CC `checkCommandOperatorPermissions()`。
+    """
+
+    command: str
+    parsed: bool
+    has_syntax_error: bool
+    permission_segments: tuple[str, ...]
+    simple_commands: tuple[BashSimpleCommand, ...]
+    has_pipeline: bool
+    has_logical_list: bool
+    has_subshell: bool
+    has_command_group: bool
+    has_command_substitution: bool
+    has_process_substitution: bool
+    has_output_redirection: bool
+    output_redirection_paths: tuple[str, ...]
+    nested_shell_command: str | None = None
+
+    @property
+    def normalized_base_commands(self) -> tuple[str, ...]:
+        return tuple(
+            cmd.normalized_base_command
+            for cmd in self.simple_commands
+            if cmd.normalized_base_command
+        )
+
+    @property
+    def has_unsafe_compound_structure(self) -> bool:
+        return self.has_subshell or self.has_command_group
+
+
+class BashCommandNameNormalizer:
+    """
+    负责将命令文本归一化到“真实基础命令”。
+
+    对齐 CC 中 wrapper stripping 的核心思想：
+    - 跳过 `FOO=1` 这类环境赋值
+    - 跳过 `env` 自身和其 flag/赋值
+    - 跳过 `command` / `builtin` / `nohup`
+    - 保留真实被执行的主命令
+    """
+
+    _ENV_ASSIGNMENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=.*$")
+    _SAFE_WRAPPERS = frozenset({"command", "builtin", "nohup", "time"})
+
+    def __init__(self, wrapper_stripper: "BashWrapperStripper | None" = None) -> None:
+        self._wrapper_stripper = wrapper_stripper or BashWrapperStripper()
+
+    def normalize(self, command_text: str) -> str | None:
+        return self.normalize_argv(self._split(command_text))
+
+    def normalize_argv(self, tokens: list[str]) -> str | None:
+        effective_tokens = list(tokens)
+        while effective_tokens and self._is_assignment(effective_tokens[0]):
+            effective_tokens = effective_tokens[1:]
+        stripped_tokens = self._wrapper_stripper.strip(effective_tokens)
+        if not stripped_tokens:
+            return None
+
+        return stripped_tokens[0].rsplit("/", 1)[-1]
+
+    def _skip_env_arguments(self, tokens: list[str], start_idx: int) -> int:
+        idx = start_idx
+        while idx < len(tokens):
+            candidate = tokens[idx]
+            if candidate == "--":
+                return idx + 1
+            if candidate.startswith("-") or self._is_assignment(candidate):
+                idx += 1
+                continue
+            return idx
+        return idx
+
+    def _is_assignment(self, token: str) -> bool:
+        return bool(self._ENV_ASSIGNMENT_RE.match(token))
+
+    @staticmethod
+    def _split(command_text: str) -> list[str]:
+        try:
+            return shlex.split(command_text)
+        except ValueError:
+            return command_text.split()
+
+
+class BashAstNodeHelper:
+    """封装 tree-sitter 节点上的常用操作。"""
+
+    def __init__(self, source: str) -> None:
+        self._source = source
+        self._source_bytes = source.encode("utf8")
+
+    def text(self, node: Node) -> str:
+        """
+        安全提取 AST 节点文本。
+
+        重要：tree-sitter 的 `start_byte/end_byte` 是按 UTF-8 字节偏移计算的，
+        不能直接对 Python `str` 做切片，否则中文等多字节字符会导致偏移错误。
+        """
+        snippet_bytes = self._source_bytes[node.start_byte:node.end_byte]
+        return snippet_bytes.decode("utf8")
+
+    def contains_node_type(self, node: Node, node_types: set[str]) -> bool:
+        if node.type in node_types:
+            return True
+        return any(self.contains_node_type(child, node_types) for child in node.children)
+
+    def contains_output_redirection(self, node: Node) -> bool:
+        if node.type == "file_redirect":
+            redirect_text = self.text(node).lstrip()
+            # `>` / `>>` / `1>` / `2>` / `&>` / `<>` 都视为写操作；纯 `<` 不算
+            if ">" in redirect_text:
+                return True
+        return any(self.contains_output_redirection(child) for child in node.children)
+
+    def command_name_text(self, node: Node) -> str | None:
+        for child in node.children:
+            if child.type == "command_name":
+                return self.text(child).strip().rsplit("/", 1)[-1]
+        return None
+
+    def argument_like_texts(self, node: Node) -> list[str]:
+        values: list[str] = []
+        for child in node.children:
+            if child.type in {
+                "command_name",
+                "word",
+                "string",
+                "raw_string",
+                "number",
+                "concatenation",
+                "expansion",
+                "simple_expansion",
+                "command_substitution",
+                "process_substitution",
+            }:
+                values.append(self.text(child).strip())
+        return [value for value in values if value]
+
+    def collect_output_redirection_paths(self, node: Node) -> list[str]:
+        paths: list[str] = []
+        self._walk_output_redirections(node, paths)
+        return paths
+
+    def _walk_output_redirections(self, node: Node, paths: list[str]) -> None:
+        if node.type == "file_redirect":
+            redirect_text = self.text(node).lstrip()
+            if ">" in redirect_text:
+                named_children = node.named_children
+                if named_children:
+                    path_text = self.text(named_children[-1]).strip()
+                    if path_text:
+                        paths.append(path_text)
+        for child in node.children:
+            self._walk_output_redirections(child, paths)
+
+
+class BashPermissionSegmenter:
+    """
+    提取用于权限检查的顶层命令段。
+
+    对齐 CC `checkCommandOperatorPermissions()`：
+    - program/list/pipeline 递归展开
+    - command/redirected_statement/subshell/compound_statement 保持原段文本
+    """
+
+    def __init__(self, node_helper: BashAstNodeHelper) -> None:
+        self._node_helper = node_helper
+
+    def collect_segments(self, node: Node) -> list[str]:
+        if node.type == "program":
+            return self._collect_from_container(node)
+
+        if node.type in {"pipeline", "list"}:
+            return self._collect_from_container(node)
+
+        if node.type in {"command", "redirected_statement", "subshell", "compound_statement"}:
+            text = self._node_helper.text(node).strip()
+            return [text] if text else []
+
+        segments = []
+        for child in node.named_children:
+            segments.extend(self.collect_segments(child))
+        if segments:
+            return [segment for segment in segments if segment.strip()]
+
+        text = self._node_helper.text(node).strip()
+        return [text] if text else []
+
+    def _collect_from_container(self, node: Node) -> list[str]:
+        segments: list[str] = []
+        for child in node.named_children:
+            segments.extend(self.collect_segments(child))
+        return [segment for segment in segments if segment.strip()]
+
+
+class BashSimpleCommandCollector:
+    """提取 AST 中全部 simple command。"""
+
+    def __init__(
+        self,
+        node_helper: BashAstNodeHelper,
+        normalizer: BashCommandNameNormalizer,
+        argv_extractor: "BashArgvExtractor | None" = None,
+    ) -> None:
+        self._node_helper = node_helper
+        self._normalizer = normalizer
+        self._argv_extractor = argv_extractor or BashArgvExtractor(node_helper)
+
+    def collect(self, node: Node) -> list[BashSimpleCommand]:
+        result: list[BashSimpleCommand] = []
+        self._walk(node, result)
+        return result
+
+    def _walk(self, node: Node, result: list[BashSimpleCommand]) -> None:
+        if node.type == "command":
+            text = self._node_helper.text(node).strip()
+            argv = tuple(self._argv_extractor.extract(node))
+            result.append(
+                BashSimpleCommand(
+                    text=text,
+                    argv=argv,
+                    base_command=self._node_helper.command_name_text(node),
+                    normalized_base_command=self._normalizer.normalize_argv(list(argv)),
+                )
+            )
+        for child in node.children:
+            self._walk(child, result)
+
+
+class BashArgvExtractor:
+    """
+    从 AST command 节点直接提取 argv。
+
+    对齐 CC validateSinglePathCommandArgv 的思路：优先使用 AST 已解析的参数，
+    避免 shell-quote / shlex 在复杂 quoting 下的歧义。
+    """
+
+    def __init__(self, node_helper: BashAstNodeHelper) -> None:
+        self._node_helper = node_helper
+
+    def extract(self, command_node: Node) -> list[str]:
+        return self._node_helper.argument_like_texts(command_node)
+
+
+class BashWrapperStripper:
+    """
+    AST argv 级 wrapper stripping。
+
+    对照 CC `stripWrappersFromArgv()`，先实现最常见且高价值的 wrappers：
+    `time` / `nohup` / `timeout` / `nice` / `stdbuf` / `env` / `command` / `builtin`
+    """
+
+    def strip(self, argv: list[str]) -> list[str]:
+        tokens = list(argv)
+        while tokens:
+            head = tokens[0]
+            if head in {"time", "nohup", "command", "builtin"}:
+                tokens = tokens[2:] if len(tokens) > 1 and tokens[1] == "--" else tokens[1:]
+                continue
+            if head == "timeout":
+                next_idx = self._skip_timeout(tokens)
+                if next_idx <= 0 or next_idx >= len(tokens):
+                    return tokens
+                tokens = tokens[next_idx + 1:]
+                continue
+            if head == "nice":
+                next_idx = self._skip_nice(tokens)
+                if next_idx <= 0 or next_idx >= len(tokens):
+                    return tokens
+                tokens = tokens[next_idx:]
+                continue
+            if head == "stdbuf":
+                next_idx = self._skip_stdbuf(tokens)
+                if next_idx <= 0 or next_idx >= len(tokens):
+                    return tokens
+                tokens = tokens[next_idx:]
+                continue
+            if head == "env":
+                next_idx = self._skip_env(tokens)
+                if next_idx <= 0 or next_idx >= len(tokens):
+                    return tokens
+                tokens = tokens[next_idx:]
+                continue
+            break
+        return tokens
+
+    def _skip_timeout(self, argv: list[str]) -> int:
+        i = 1
+        while i < len(argv):
+            arg = argv[i]
+            next_arg = argv[i + 1] if i + 1 < len(argv) else None
+            if arg in {"--foreground", "--preserve-status", "--verbose"}:
+                i += 1
+            elif arg.startswith("--kill-after=") or arg.startswith("--signal="):
+                i += 1
+            elif arg in {"--kill-after", "--signal"} and next_arg:
+                i += 2
+            elif arg == "--":
+                i += 1
+                break
+            elif arg.startswith("--"):
+                return -1
+            elif arg == "-v":
+                i += 1
+            elif arg in {"-k", "-s"} and next_arg:
+                i += 2
+            elif re.match(r"^-[ks][A-Za-z0-9_.+-]+$", arg):
+                i += 1
+            elif arg.startswith("-"):
+                return -1
+            else:
+                break
+        return i
+
+    def _skip_nice(self, argv: list[str]) -> int:
+        if len(argv) >= 3 and argv[1] == "-n" and re.match(r"^-?\d+$", argv[2]):
+            return 4 if len(argv) > 3 and argv[3] == "--" else 3
+        if len(argv) >= 2 and re.match(r"^-\d+$", argv[1]):
+            return 3 if len(argv) > 2 and argv[2] == "--" else 2
+        return 2 if len(argv) > 1 and argv[1] == "--" else 1
+
+    def _skip_stdbuf(self, argv: list[str]) -> int:
+        i = 1
+        while i < len(argv):
+            arg = argv[i]
+            if re.match(r"^-[ioe]$", arg) and i + 1 < len(argv):
+                i += 2
+            elif re.match(r"^-[ioe].", arg):
+                i += 1
+            elif re.match(r"^--(input|output|error)=", arg):
+                i += 1
+            elif arg.startswith("-"):
+                return -1
+            else:
+                break
+        return i if i > 1 and i < len(argv) else -1
+
+    def _skip_env(self, argv: list[str]) -> int:
+        i = 1
+        while i < len(argv):
+            arg = argv[i]
+            if "=" in arg and not arg.startswith("-"):
+                i += 1
+            elif arg in {"-i", "-0", "-v"}:
+                i += 1
+            elif arg == "-u" and i + 1 < len(argv):
+                i += 2
+            elif arg.startswith("-"):
+                return -1
+            else:
+                break
+        return i if i < len(argv) else -1
+
+
+class BashAstAdvisor:
+    """
+    基于 AST 分析结果给出权限建议。
+
+    对齐 CC 的思路：复杂 shell 结构优先 ask，而不是盲目依赖字符串规则。
+    """
+
+    def get_approval_reason(self, analysis: BashAstAnalysis) -> str | None:
+        if analysis.has_subshell:
+            return "This command uses a subshell `(...)`, which requires approval for safety."
+        if analysis.has_command_group:
+            return "This command uses a command group `{ ...; }`, which requires approval for safety."
+        if analysis.has_command_substitution:
+            return "This command uses command substitution (`$()` or backticks), which requires approval for safety."
+        if analysis.has_process_substitution:
+            return "This command uses process substitution (`<()` or `>()`), which requires approval for safety."
+        if analysis.nested_shell_command:
+            return (
+                f"This command launches a nested shell via `{analysis.nested_shell_command} -c`, "
+                "which requires approval for safety."
+            )
+        if analysis.has_syntax_error:
+            return "This command has shell syntax errors and requires approval before execution."
+
+        normalized = analysis.normalized_base_commands
+        if sum(1 for command in normalized if command == "cd") > 1:
+            return "Multiple directory changes in one command require approval for clarity."
+
+        if "cd" in normalized and "git" in normalized and len(analysis.permission_segments) > 1:
+            return (
+                "Compound commands mixing `cd` and `git` across segments require approval "
+                "to avoid repository-boundary confusion."
+            )
+
+        return None
+
+    def requires_write_permissions(self, analysis: BashAstAnalysis) -> bool:
+        return analysis.has_output_redirection
+
+
+class BashAstAnalyzer:
+    """
+    Bash AST 主分析器。
+
+    这是 OOP 入口，供 `BashRuntimeTool` 调用。
+    """
+
+    _SHELL_WRAPPER_COMMANDS = frozenset({"bash", "sh", "zsh", "fish"})
+
+    def __init__(self) -> None:
+        self._normalizer = BashCommandNameNormalizer()
+        self._advisor = BashAstAdvisor()
+
+    def analyze(self, command: str) -> BashAstAnalysis:
+        stripped = command.strip()
+        if not stripped:
+            return BashAstAnalysis(
+                command=command,
+                parsed=True,
+                has_syntax_error=False,
+                permission_segments=(),
+                simple_commands=(),
+                has_pipeline=False,
+                has_logical_list=False,
+                has_subshell=False,
+                has_command_group=False,
+                has_command_substitution=False,
+                has_process_substitution=False,
+                has_output_redirection=False,
+                output_redirection_paths=(),
+            )
+
+        try:
+            parser = self._get_bash_parser()
+            tree = parser.parse(command.encode("utf-8"))
+            root = tree.root_node
+        except Exception:
+            return self._build_fallback_analysis(command)
+
+        node_helper = BashAstNodeHelper(command)
+        segmenter = BashPermissionSegmenter(node_helper)
+        collector = BashSimpleCommandCollector(node_helper, self._normalizer)
+
+        simple_commands = tuple(collector.collect(root))
+        permission_segments = tuple(segmenter.collect_segments(root))
+
+        return BashAstAnalysis(
+            command=command,
+            parsed=True,
+            has_syntax_error=root.has_error,
+            permission_segments=permission_segments or ((stripped,) if stripped else ()),
+            simple_commands=simple_commands,
+            has_pipeline=node_helper.contains_node_type(root, {"pipeline"}),
+            has_logical_list=node_helper.contains_node_type(root, {"list"}),
+            has_subshell=node_helper.contains_node_type(root, {"subshell"}),
+            has_command_group=node_helper.contains_node_type(root, {"compound_statement"}),
+            has_command_substitution=node_helper.contains_node_type(root, {"command_substitution"}),
+            has_process_substitution=node_helper.contains_node_type(root, {"process_substitution"}),
+            has_output_redirection=node_helper.contains_output_redirection(root),
+            output_redirection_paths=tuple(node_helper.collect_output_redirection_paths(root)),
+            nested_shell_command=self._detect_nested_shell(simple_commands),
+        )
+
+    def get_approval_reason(self, analysis: BashAstAnalysis) -> str | None:
+        return self._advisor.get_approval_reason(analysis)
+
+    def requires_write_permissions(self, analysis: BashAstAnalysis) -> bool:
+        return self._advisor.requires_write_permissions(analysis)
+
+    def _build_fallback_analysis(self, command: str) -> BashAstAnalysis:
+        stripped = command.strip()
+        return BashAstAnalysis(
+            command=command,
+            parsed=False,
+            has_syntax_error=False,
+            permission_segments=(stripped,) if stripped else (),
+            simple_commands=(),
+            has_pipeline="|" in stripped,
+            has_logical_list="&&" in stripped or "||" in stripped or ";" in stripped,
+            has_subshell=False,
+            has_command_group=False,
+            has_command_substitution=False,
+            has_process_substitution=False,
+            has_output_redirection=">" in stripped,
+            output_redirection_paths=(),
+        )
+
+    def _detect_nested_shell(
+        self,
+        simple_commands: tuple[BashSimpleCommand, ...],
+    ) -> str | None:
+        for command in simple_commands:
+            normalized = command.normalized_base_command
+            if normalized not in self._SHELL_WRAPPER_COMMANDS:
+                continue
+            tokens = self._safe_split(command.text)
+            if "-c" in tokens:
+                return normalized
+        return None
+
+    @staticmethod
+    def _safe_split(command_text: str) -> list[str]:
+        try:
+            return shlex.split(command_text)
+        except ValueError:
+            return command_text.split()
+
+    @staticmethod
+    @lru_cache(maxsize=1)
+    def _get_bash_parser():
+        return get_parser("bash")
+
+
+_DEFAULT_ANALYZER = BashAstAnalyzer()
+
+
+def analyze_command_ast(command: str) -> BashAstAnalysis:
+    """兼容入口：分析 bash 命令 AST。"""
+    return _DEFAULT_ANALYZER.analyze(command)
+
+
+def get_ast_approval_reason(analysis: BashAstAnalysis) -> str | None:
+    """兼容入口：返回 AST 级审批原因。"""
+    return _DEFAULT_ANALYZER.get_approval_reason(analysis)
+
+
+def command_requires_write_permissions(command: str) -> bool:
+    """兼容入口：判断命令是否需要写权限。"""
+    return _DEFAULT_ANALYZER.requires_write_permissions(analyze_command_ast(command))