tools: enforce BashTool whitelist across compound commands

tests/file_tools/test_bash_tool.py

@@ -6,6 +6,7 @@
 """Unit tests for BashTool."""
 
 from unittest.mock import Mock
+from unittest.mock import patch
 
 import pytest
 from trpc_agent_sdk.context import InvocationContext
@@ -183,8 +184,121 @@ def test_is_command_safe(self, tool, tmp_path):
         assert tool._is_command_safe("echo test", str(tmp_path)) is True
         assert tool._is_command_safe("ls", "/tmp") is True
 
-        import os
+        assert tool._is_command_safe("blocked_cmd", str(tmp_path)) is True
+
         workdir = str(tmp_path)
         outside_dir = "/var" if workdir != "/var" else "/usr"
-        result = tool._is_command_safe("rm -rf /nonexistent", outside_dir)
-        assert isinstance(result, bool)
+        assert tool._is_command_safe("blocked_cmd", outside_dir) is False
+
+    @pytest.mark.parametrize(
+        "command",
+        [
+            "blocked_cmd",
+            "echo ok | blocked_cmd",
+            "echo ok; blocked_cmd",
+            "echo ok && blocked_cmd",
+            "echo ok || blocked_cmd",
+            "echo ok\nblocked_cmd",
+            "echo ok & blocked_cmd",
+            "echo $(blocked_cmd)",
+            "echo `blocked_cmd`",
+            r"echo escaped\>& blocked_cmd",
+        ],
+    )
+    def test_is_command_safe_rejects_non_whitelisted_command_segments(self, tool_with_whitelist, command):
+        """Every executable command segment must be allowlisted."""
+        assert tool_with_whitelist._is_command_safe(command, tool_with_whitelist.cwd) is False
+
+    @pytest.mark.parametrize(
+        "command",
+        [
+            "echo ok",
+            "echo ok | echo done",
+            "echo ok; echo done",
+            "echo ok && echo done",
+            "echo ok || echo done",
+            "echo ok\necho done",
+            "echo ok & echo done",
+        ],
+    )
+    def test_is_command_safe_allows_only_whitelisted_command_segments(self, tool_with_whitelist, command):
+        """Compound commands remain valid when every command is allowlisted."""
+        assert tool_with_whitelist._is_command_safe(command, tool_with_whitelist.cwd) is True
+
+    @pytest.mark.parametrize(
+        "command",
+        [
+            "echo 'literal; not a separator'",
+            'echo "literal && not a separator"',
+            r"echo escaped\;separator",
+            "echo ok 2>&1",
+            "echo ok &>output.log",
+            "echo '$" + "\\" + "\n" + "(blocked_cmd)'",
+        ],
+    )
+    def test_is_command_safe_preserves_quoted_escaped_and_redirected_arguments(self, tool_with_whitelist, command):
+        """Shell syntax inside arguments and file-descriptor redirects is not a command boundary."""
+        assert tool_with_whitelist._is_command_safe(command, tool_with_whitelist.cwd) is True
+
+    @pytest.mark.parametrize(
+        "command",
+        [
+            "cat <<EOF\nplain data\nEOF",
+            "cat <<'EOF'\n$(blocked_cmd)\nEOF",
+            "cat <<-EOF\n\ttab-indented data\n\tEOF",
+            "cat <<EOF\nplain data\nEOF\necho done",
+        ],
+    )
+    def test_is_command_safe_allows_heredoc_data(self, tmp_path, command):
+        """Heredoc bodies are data rather than separate command segments."""
+        tool = BashTool(cwd=str(tmp_path), whitelist_commands=["cat", "echo"])
+
+        assert tool._is_command_safe(command, tool.cwd) is True
+
+    @pytest.mark.parametrize(
+        "command",
+        [
+            "cat <<EOF\nplain data",
+            "cat <<EOF\n$(blocked_cmd)\nEOF",
+            "cat <<EOF\nplain data\nEOF\nblocked_cmd",
+        ],
+    )
+    def test_is_command_safe_rejects_unsafe_or_unterminated_heredocs(self, tmp_path, command):
+        """Heredocs fail closed when parsing or expansion is unsafe."""
+        tool = BashTool(cwd=str(tmp_path), whitelist_commands=["cat", "echo"])
+
+        assert tool._is_command_safe(command, tool.cwd) is False
+
+    @pytest.mark.parametrize(
+        "command",
+        [
+            "echo 'unterminated",
+            "echo <(blocked_cmd)",
+            "echo >(blocked_cmd)",
+            "echo $" + "\\" + "\n" + "(blocked_cmd)",
+            'echo "$' + "\\" + "\n" + '(blocked_cmd)"',
+            "cat <" + "\\" + "\n" + "(blocked_cmd)",
+            "cat >" + "\\" + "\n" + "(blocked_cmd)",
+        ],
+    )
+    def test_is_command_safe_fails_closed_for_unverifiable_syntax(self, tool_with_whitelist, command):
+        """Unparseable syntax and process substitution are rejected."""
+        assert tool_with_whitelist._is_command_safe(command, tool_with_whitelist.cwd) is False
+
+    @pytest.mark.asyncio
+    @patch("trpc_agent_sdk.tools.file_tools._bash_tool.asyncio.create_subprocess_shell")
+    async def test_bash_custom_whitelist_blocks_before_subprocess(
+        self,
+        mock_create_subprocess_shell,
+        tool_with_whitelist,
+        tool_context,
+    ):
+        """A rejected command never reaches the subprocess boundary."""
+        result = await tool_with_whitelist._run_async_impl(
+            tool_context=tool_context,
+            args={"command": "blocked_cmd"},
+        )
+
+        assert result["success"] is False
+        assert "SECURITY_RESTRICTION" in result["error"]
+        mock_create_subprocess_shell.assert_not_called()

trpc_agent_sdk/tools/file_tools/_bash_tool.py

@@ -27,6 +27,199 @@
 from trpc_agent_sdk.types import Type
 
 
+def _parse_heredoc_operator(command: str, index: int) -> Optional[tuple[int, str, bool, bool]]:
+    """Parse a heredoc operator and its delimiter without running a shell."""
+    cursor = index + 2
+    strip_tabs = cursor < len(command) and command[cursor] == "-"
+    if strip_tabs:
+        cursor += 1
+
+    while cursor < len(command) and command[cursor] in {" ", "\t"}:
+        cursor += 1
+
+    delimiter: list[str] = []
+    quote = ""
+    quoted = False
+    while cursor < len(command):
+        char = command[cursor]
+        if quote:
+            if char == quote:
+                quote = ""
+            elif char == "\\" and quote == '"':
+                quoted = True
+                cursor += 1
+                if cursor >= len(command) or command[cursor] == "\n":
+                    return None
+                delimiter.append(command[cursor])
+            else:
+                delimiter.append(char)
+            cursor += 1
+            continue
+
+        if char in {"'", '"'}:
+            quote = char
+            quoted = True
+            cursor += 1
+            continue
+        if char == "\\":
+            quoted = True
+            cursor += 1
+            if cursor >= len(command) or command[cursor] == "\n":
+                return None
+            delimiter.append(command[cursor])
+            cursor += 1
+            continue
+        if char.isspace() or char in {"|", ";", "&", "<", ">"}:
+            break
+        delimiter.append(char)
+        cursor += 1
+
+    if quote or not delimiter:
+        return None
+    return cursor, "".join(delimiter), strip_tabs, quoted
+
+
+def _skip_heredoc_bodies(
+    command: str,
+    index: int,
+    heredocs: list[tuple[str, bool, bool]],
+) -> Optional[int]:
+    """Skip heredoc data while rejecting executable substitutions."""
+    cursor = index
+    for delimiter, strip_tabs, quoted in heredocs:
+        while cursor <= len(command):
+            line_end = command.find("\n", cursor)
+            if line_end == -1:
+                line_end = len(command)
+            line = command[cursor:line_end].removesuffix("\r")
+            comparable_line = line.lstrip("\t") if strip_tabs else line
+            if comparable_line == delimiter:
+                cursor = line_end + 1 if line_end < len(command) else line_end
+                break
+            if not quoted and any(pattern in line for pattern in ("$(", "`", "<(", ">(")):
+                return None
+            if line_end == len(command):
+                return None
+            cursor = line_end + 1
+        else:
+            return None
+    return cursor
+
+
+def _extract_base_commands(command: str) -> Optional[list[str]]:
+    """Extract every executable command segment without executing the shell."""
+    segments: list[str] = []
+    current: list[str] = []
+    quote = ""
+    escaped = False
+    index = 0
+    previous_is_redirect = False
+    pending_heredocs: list[tuple[str, bool, bool]] = []
+
+    while index < len(command):
+        char = command[index]
+
+        if escaped:
+            current.append(char)
+            escaped = False
+            previous_is_redirect = False
+            index += 1
+            continue
+
+        if char == "\\" and quote != "'":
+            if command.startswith(("\\\n", "\\\r\n"), index):
+                return None
+            current.append(char)
+            escaped = True
+            previous_is_redirect = False
+            index += 1
+            continue
+
+        if quote:
+            current.append(char)
+            if char == quote:
+                quote = ""
+            elif quote == '"' and (char == "`" or command.startswith("$(", index)):
+                return None
+            previous_is_redirect = False
+            index += 1
+            continue
+
+        if char in {"'", '"'}:
+            quote = char
+            current.append(char)
+            previous_is_redirect = False
+            index += 1
+            continue
+
+        if char == "`" or command.startswith(("$(", "<(", ">("), index):
+            return None
+
+        if command.startswith("<<", index) and not command.startswith("<<<", index):
+            operator_start = index
+            heredoc = _parse_heredoc_operator(command, index)
+            if heredoc is None:
+                return None
+            index, delimiter, strip_tabs, quoted = heredoc
+            current.append(command[operator_start:index])
+            pending_heredocs.append((delimiter, strip_tabs, quoted))
+            previous_is_redirect = False
+            continue
+
+        if char in {"|", ";", "&", "\n"}:
+            # Keep file-descriptor redirections such as 2>&1 and &>file in
+            # the current segment; a standalone ampersand is a separator.
+            next_is_redirect = index + 1 < len(command) and command[index + 1] == ">"
+            is_redirection = char == "&" and (previous_is_redirect or next_is_redirect)
+            if is_redirection:
+                current.append(char)
+                previous_is_redirect = False
+                index += 1
+                continue
+
+            segment = "".join(current).strip()
+            if segment:
+                segments.append(segment)
+            current = []
+            previous_is_redirect = False
+            if char == "\n" and pending_heredocs:
+                heredoc_end = _skip_heredoc_bodies(command, index + 1, pending_heredocs)
+                if heredoc_end is None:
+                    return None
+                pending_heredocs = []
+                index = heredoc_end
+                continue
+            if command.startswith(("||", "&&", "|&", ";;"), index):
+                index += 2
+            else:
+                index += 1
+            continue
+
+        current.append(char)
+        previous_is_redirect = char in {"<", ">"}
+        index += 1
+
+    if quote or escaped or pending_heredocs:
+        return None
+
+    segment = "".join(current).strip()
+    if segment:
+        segments.append(segment)
+    if not segments:
+        return None
+
+    base_commands: list[str] = []
+    for segment in segments:
+        try:
+            tokens = shlex.split(segment, posix=True)
+        except ValueError:
+            return None
+        if not tokens:
+            return None
+        base_commands.append(tokens[0])
+    return base_commands
+
+
 class BashTool(BaseTool):
     """Tool for executing bash commands."""
 
@@ -127,36 +320,20 @@ def _is_command_safe(self, command: str, execution_dir: str) -> bool:
         Returns:
             Whether it's safe to execute
         """
-        try:
-            lexer = shlex.shlex(command, posix=True, punctuation_chars="|")
-            lexer.whitespace_split = True
-            tokens = list(lexer)
-
-            base_commands = []
-            current_command = []
-            for token in tokens:
-                if token == "|":
-                    if current_command:
-                        base_commands.append(current_command[0])
-                        current_command = []
-                else:
-                    if not current_command:
-                        current_command.append(token)
-        except Exception:  # pylint: disable=broad-except
-            base_commands = [command.split()[0] if command.split() else ""]
-
-        for base_command in base_commands:
-            if self.whitelist_commands is not None:
-                if base_command not in self.whitelist_commands:
-                    return False
-            else:
-                try:
-                    Path(execution_dir).resolve().relative_to(Path(self.cwd).resolve())
-                except ValueError:
-                    if base_command not in self.ALLOWED_COMMANDS_OUTSIDE_WORKDIR:
-                        return False
+        if self.whitelist_commands is not None:
+            allowed_commands = self.whitelist_commands
+        else:
+            try:
+                Path(execution_dir).resolve().relative_to(Path(self.cwd).resolve())
+                return True
+            except ValueError:
+                allowed_commands = self.ALLOWED_COMMANDS_OUTSIDE_WORKDIR
+
+        base_commands = _extract_base_commands(command)
+        if not base_commands:
+            return False
 
-        return True
+        return all(base_command in allowed_commands for base_command in base_commands)
 
     async def _run_async_impl(self, *, tool_context: InvocationContext, args: dict[str, Any]) -> Any:
         command = args.get("command")