From 2d1ea7bc291c21cc80f046e37d0e2f079344103a Mon Sep 17 00:00:00 2001
From: notgitika <gitijh@gmail.com>
Date: Thu, 16 Jul 2026 18:03:51 -0400
Subject: [PATCH] fix: tighten package specifier validation in
 install_packages()

Restrict the extras group in the package-name pattern to valid extra
names instead of allowing arbitrary characters, and quote each package
argument with shlex.quote when building the pip command.

Adds tests for extras handling and argument quoting.
---
 .../tools/code_interpreter_client.py          |  9 +++-
 .../tools/test_code_interpreter_client.py     | 50 ++++++++++++++++++-
 2 files changed, 55 insertions(+), 4 deletions(-)

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