diff --git a/.changes/next-version.json b/.changes/next-version.json new file mode 100644 index 0000000000..0f15dd2a69 --- /dev/null +++ b/.changes/next-version.json @@ -0,0 +1,7 @@ +[ + { + "category": "Session", + "description": "Add support for ``aws_bearer_token`` parameter in Session, Client, and Resource to enable per-session/client bearer token authentication for services that support it (e.g. Bedrock). This allows safe multi-tenant usage without relying on global environment variables like ``AWS_BEARER_TOKEN_BEDROCK``.", + "type": "feature" + } +] diff --git a/boto3/session.py b/boto3/session.py index df978595ae..d96af52607 100644 --- a/boto3/session.py +++ b/boto3/session.py @@ -50,6 +50,9 @@ class Session: the default profile is used. :type aws_account_id: string :param aws_account_id: AWS account ID + :type aws_bearer_token: string + :param aws_bearer_token: Bearer token for authentication with services + that support bearer token auth (e.g. Bedrock). """ def __init__( @@ -61,6 +64,7 @@ def __init__( botocore_session=None, profile_name=None, aws_account_id=None, + aws_bearer_token=None, ): if botocore_session is not None: self._session = botocore_session @@ -100,6 +104,9 @@ def __init__( if region_name is not None: self._session.set_config_variable('region', region_name) + # Store the bearer token for per-session bearer token authentication + self._aws_bearer_token = aws_bearer_token + self.resource_factory = ResourceFactory( self._session.get_component('event_emitter') ) @@ -243,6 +250,7 @@ def client( aws_session_token=None, config=None, aws_account_id=None, + aws_bearer_token=None, ): """ Create a low-level service client by name. @@ -314,9 +322,25 @@ def client( :param aws_account_id: The account id to use when creating the client. Same semantics as aws_access_key_id above. + :type aws_bearer_token: string + :param aws_bearer_token: The bearer token to use for + authentication with services that support bearer token auth + (e.g. Bedrock). If provided, this takes precedence over any + bearer token resolved from environment variables and any + session-level token. Same semantics as aws_access_key_id + above. + :return: Service client instance """ + # Determine the effective bearer token to use + # Precedence: client parameter > session parameter > environment variable + effective_aws_bearer_token = ( + aws_bearer_token + if aws_bearer_token is not None + else self._aws_bearer_token + ) + create_client_kwargs = { 'region_name': region_name, 'api_version': api_version, @@ -328,11 +352,16 @@ def client( 'aws_session_token': aws_session_token, 'config': config, 'aws_account_id': aws_account_id, + 'aws_bearer_token': effective_aws_bearer_token, } if aws_account_id is None: # Remove aws_account_id for arbitrary # botocore version mismatches in AWS Lambda. del create_client_kwargs['aws_account_id'] + if effective_aws_bearer_token is None: + # Remove aws_bearer_token for arbitrary + # botocore version mismatches in AWS Lambda. + del create_client_kwargs['aws_bearer_token'] return self._session.create_client( service_name, **create_client_kwargs @@ -350,6 +379,7 @@ def resource( aws_secret_access_key=None, aws_session_token=None, config=None, + aws_bearer_token=None, ): """ Create a resource service client by name. @@ -419,6 +449,14 @@ def resource( `_ for more details. + :type aws_bearer_token: string + :param aws_bearer_token: The bearer token to use for + authentication with services that support bearer token auth + (e.g. Bedrock). If provided, this takes precedence over any + bearer token resolved from environment variables and any + session-level token. Same semantics as aws_access_key_id + above. + :return: Subclass of :py:class:`~boto3.resources.base.ServiceResource` """ try: @@ -483,6 +521,7 @@ def resource( aws_secret_access_key=aws_secret_access_key, aws_session_token=aws_session_token, config=config, + aws_bearer_token=aws_bearer_token, ) service_model = client.meta.service_model diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index f189b87948..8a3638c2dd 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -272,6 +272,116 @@ def test_create_client_with_args(self): config=None, ) + def test_session_with_aws_bearer_token(self): + session = Session(aws_bearer_token='test-token') + assert session._aws_bearer_token == 'test-token' + + def test_session_with_aws_bearer_token_none(self): + session = Session() + assert session._aws_bearer_token is None + + def test_create_client_with_aws_bearer_token(self): + bc_session = self.bc_session_cls.return_value + + session = Session(region_name='us-east-1') + session.client( + 'bedrock-runtime', + region_name='us-west-2', + aws_bearer_token='test-bearer-token', + ) + + bc_session.create_client.assert_called_with( + 'bedrock-runtime', + aws_secret_access_key=None, + aws_access_key_id=None, + endpoint_url=None, + use_ssl=True, + aws_session_token=None, + verify=None, + region_name='us-west-2', + api_version=None, + config=None, + aws_bearer_token='test-bearer-token', + ) + + def test_create_client_with_session_level_aws_bearer_token(self): + bc_session = self.bc_session_cls.return_value + + session = Session( + region_name='us-east-1', + aws_bearer_token='session-level-token', + ) + session.client('bedrock-runtime', region_name='us-west-2') + + bc_session.create_client.assert_called_with( + 'bedrock-runtime', + aws_secret_access_key=None, + aws_access_key_id=None, + endpoint_url=None, + use_ssl=True, + aws_session_token=None, + verify=None, + region_name='us-west-2', + api_version=None, + config=None, + aws_bearer_token='session-level-token', + ) + + def test_create_client_aws_bearer_token_client_overrides_session(self): + bc_session = self.bc_session_cls.return_value + + session = Session( + region_name='us-east-1', + aws_bearer_token='session-level-token', + ) + session.client( + 'bedrock-runtime', + region_name='us-west-2', + aws_bearer_token='client-level-token', + ) + + bc_session.create_client.assert_called_with( + 'bedrock-runtime', + aws_secret_access_key=None, + aws_access_key_id=None, + endpoint_url=None, + use_ssl=True, + aws_session_token=None, + verify=None, + region_name='us-west-2', + api_version=None, + config=None, + aws_bearer_token='client-level-token', + ) + + def test_create_client_aws_bearer_token_empty_string_overrides_session(self): + bc_session = self.bc_session_cls.return_value + + session = Session( + region_name='us-east-1', + aws_bearer_token='session-level-token', + ) + # Empty string is an explicit value and should override session token + session.client( + 'bedrock-runtime', + region_name='us-west-2', + aws_bearer_token='', + ) + + bc_session.create_client.assert_called_with( + 'bedrock-runtime', + aws_secret_access_key=None, + aws_access_key_id=None, + endpoint_url=None, + use_ssl=True, + aws_session_token=None, + verify=None, + region_name='us-west-2', + api_version=None, + config=None, + aws_bearer_token='', + ) + def test_create_client_with_aws_account_id(self): bc_session = self.bc_session_cls.return_value @@ -324,6 +434,7 @@ def test_create_resource_with_args(self): region_name=None, api_version='2014-11-02', config=mock.ANY, + aws_bearer_token=None, ) client_config = session.client.call_args[1]['config'] assert client_config.user_agent_extra == 'Resource' @@ -356,6 +467,7 @@ def test_create_resource_with_config(self): region_name=None, api_version='2014-11-02', config=mock.ANY, + aws_bearer_token=None, ) client_config = session.client.call_args[1]['config'] assert client_config.user_agent_extra == 'Resource' @@ -388,6 +500,7 @@ def test_create_resource_with_config_override_user_agent_extra(self): region_name=None, api_version='2014-11-02', config=mock.ANY, + aws_bearer_token=None, ) client_config = session.client.call_args[1]['config'] assert client_config.user_agent_extra == 'foo'