diff --git a/tests/file_tools/test_bash_tool.py b/tests/file_tools/test_bash_tool.py index bab9c0aa7..40155929c 100644 --- a/tests/file_tools/test_bash_tool.py +++ b/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 <(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() diff --git a/trpc_agent_sdk/tools/file_tools/_bash_tool.py b/trpc_agent_sdk/tools/file_tools/_bash_tool.py index 6386e9e4c..1d3f63d06 100644 --- a/trpc_agent_sdk/tools/file_tools/_bash_tool.py +++ b/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")