diff --git a/src/bedrock_agentcore/tools/code_interpreter_client.py b/src/bedrock_agentcore/tools/code_interpreter_client.py index 147f8d08..0cc98484 100644 --- a/src/bedrock_agentcore/tools/code_interpreter_client.py +++ b/src/bedrock_agentcore/tools/code_interpreter_client.py @@ -6,6 +6,7 @@ import logging import re +import shlex import uuid from contextlib import contextmanager from typing import Any, Dict, Generator, List, Optional, Union @@ -20,8 +21,11 @@ DEFAULT_IDENTIFIER = "aws.codeinterpreter.v1" +# Allowlist for pip package specifiers. The extras group is restricted to valid +# extra names (comma-separated identifiers) rather than allowing arbitrary +# characters, so only well-formed package specifiers are accepted. VALID_PACKAGE_NAME = re.compile( - r"^[a-zA-Z0-9]([a-zA-Z0-9._-]*[a-zA-Z0-9])?(\[.*\])?(==|>=|<=|!=|~=|>|<)?[a-zA-Z0-9.*]*$" + r"^[a-zA-Z0-9]([a-zA-Z0-9._-]*[a-zA-Z0-9])?(\[[a-zA-Z0-9._,\-]*\])?(==|>=|<=|!=|~=|>|<)?[a-zA-Z0-9.*]*$" ) DEFAULT_TIMEOUT = 900 @@ -606,7 +610,8 @@ def install_packages( if not VALID_PACKAGE_NAME.match(pkg): raise ValueError(f"Invalid package name: {pkg}") - packages_str = " ".join(packages) + # Quote each argument so it is passed to the command as a single literal token. + packages_str = " ".join(shlex.quote(pkg) for pkg in packages) upgrade_flag = "--upgrade " if upgrade else "" command = f"pip install {upgrade_flag}{packages_str}" diff --git a/tests/bedrock_agentcore/tools/test_code_interpreter_client.py b/tests/bedrock_agentcore/tools/test_code_interpreter_client.py index 8c968481..a3430957 100644 --- a/tests/bedrock_agentcore/tools/test_code_interpreter_client.py +++ b/tests/bedrock_agentcore/tools/test_code_interpreter_client.py @@ -838,12 +838,12 @@ def test_install_packages_with_versions(self, mock_boto3): # Act client.install_packages(["pandas>=2.0", "numpy<2.0", "scikit-learn==1.3.0"]) - # Assert + # Assert - version specifiers are shell-quoted (harmless; pip sees the original) client.data_plane_client.invoke_code_interpreter.assert_called_once_with( codeInterpreterIdentifier="test.identifier", sessionId="test-session-id", name="executeCommand", - arguments={"command": "pip install pandas>=2.0 numpy<2.0 scikit-learn==1.3.0"}, + arguments={"command": "pip install 'pandas>=2.0' 'numpy<2.0' scikit-learn==1.3.0"}, ) @patch("bedrock_agentcore.tools.code_interpreter_client.boto3") @@ -914,6 +914,52 @@ def test_install_packages_invalid_characters_error(self, mock_boto3): with pytest.raises(ValueError, match="Invalid package name"): client.install_packages(["pandas$HOME"]) + @patch("bedrock_agentcore.tools.code_interpreter_client.boto3") + def test_install_packages_extras_group_metacharacters_rejected(self, mock_boto3): + """Extras groups containing characters outside valid extra names are rejected.""" + # Arrange + mock_session = MagicMock() + mock_session.client.return_value = MagicMock() + mock_boto3.Session.return_value = mock_session + client = CodeInterpreter("us-west-2") + client.identifier = "test.identifier" + client.session_id = "test-session-id" + + # Act & Assert - metacharacters inside an extras group are not valid extra names + with pytest.raises(ValueError, match="Invalid package name"): + client.install_packages(["requests[$(id)]"]) + + with pytest.raises(ValueError, match="Invalid package name"): + client.install_packages(["flask[`whoami`]"]) + + with pytest.raises(ValueError, match="Invalid package name"): + client.install_packages(["pandas[a b c]"]) + + @patch("bedrock_agentcore.tools.code_interpreter_client.boto3") + def test_install_packages_valid_extras_allowed(self, mock_boto3): + """Legitimate pip extras (comma-separated identifiers) must still work.""" + # Arrange + mock_session = MagicMock() + mock_session.client.return_value = MagicMock() + mock_boto3.Session.return_value = mock_session + client = CodeInterpreter("us-west-2") + client.identifier = "test.identifier" + client.session_id = "test-session-id" + + mock_response = {"result": "success"} + client.data_plane_client.invoke_code_interpreter.return_value = mock_response + + # Act + client.install_packages(["pandas[security]", "celery[redis,auth]"]) + + # Assert + client.data_plane_client.invoke_code_interpreter.assert_called_once_with( + codeInterpreterIdentifier="test.identifier", + sessionId="test-session-id", + name="executeCommand", + arguments={"command": "pip install 'pandas[security]' 'celery[redis,auth]'"}, + ) + @patch("bedrock_agentcore.tools.code_interpreter_client.boto3") def test_download_file_text(self, mock_boto3): # Arrange