Assume Role Client

credential_providers.py

from botocore.session import Session
from botocore.credentials import AssumeRoleProvider, CredentialResolver, CanonicalNameCredentialSourcer

"""
# Where much of this is referenced:
https://github.com/aws/aws-cli/blob/f4ea4b682e792eb474e9e3c3a74b16ab34ba95c0/tests/integration/test_cli.py
https://github.com/boto/botocore/blob/fd6aac245b24dbdbfc6e0c5ecd7530cfe4346794/botocore/session.py
https://github.com/boto/botocore/blob/fd6aac245b24dbdbfc6e0c5ecd7530cfe4346794/botocore/credentials.py
"""


def get_bucket_snake_botocore_session(source_app, dest_account, session_name="bucketsnake", **botocore_session_kwargs):
    """
    This will create a new Bucket Snake botocore session -- under the hood this leverages the
    AssumeRoleCredentialProvider. This works by faking a profile and entering in the correct role arn.
    :param source_app:
    :param dest_account:
    :param session_name:
    :param botocore_session_kwargs:
    :return:
    """
    session = Session(**botocore_session_kwargs)

    # Get the current source credential provider:
    source_provider = session._components.get_component("credential_provider")

    # Get the source provider name -- this is needed by the AssumeRoleCredentialProviders (roundabout but it works):
    spn = source_provider.get_provider(source_provider.load_credentials().method).CANONICAL_NAME
    
    # TODO: Get the source AWS account number... sts-get-caller-identity??
    
    # Fake the configuration:
    profile_config = {"profiles": {
        "bucketsnake": {
            "role_arn": "arn:aws:iam::{account}:role/{app}-{account}".format(account=dest_account, app=source_app),
            "role_session_name": session_name,
            "credential_source": spn
        }
    }}
    session.full_config.update(profile_config)

    # Overwrite the credential resolver to our own:
    session._components.lazy_register_component(
        'credential_provider',
        lambda: create_bucket_snake_credential_resolver(session, source_provider))

    return session


def create_bucket_snake_credential_resolver(session, source_provider):
    """Create the Bucket Snake credential resolver with the required parameters"""
    bucketsnake_provider = AssumeRoleProvider(
        load_config=lambda: session.full_config,
        client_creator=session.create_client,
        cache={},
        profile_name="bucketsnake",
        credential_sourcer=CanonicalNameCredentialSourcer(source_provider.providers)
    )
    bucketsnake_provider.METHOD = "bucket-snake"

    resolver = CredentialResolver(providers=[bucketsnake_provider])

    return resolver

OLD

# from botocore.session import Session
# from botocore.credentials import AssumeRoleProvider, CredentialResolver, AssumeRoleCredentialFetcher, \
#     DeferredRefreshableCredentials, _local_now

# """
# # Where much of this is referenced:
# https://github.com/aws/aws-cli/blob/f4ea4b682e792eb474e9e3c3a74b16ab34ba95c0/tests/integration/test_cli.py
# https://github.com/boto/botocore/blob/fd6aac245b24dbdbfc6e0c5ecd7530cfe4346794/botocore/session.py
# https://github.com/boto/botocore/blob/fd6aac245b24dbdbfc6e0c5ecd7530cfe4346794/botocore/credentials.py
# """


# def create_bucket_snake_credential_resolver(session, source_app, dest_account, role_session_name, source_provider):
#     """Create a default credential resolver.
#     This creates a pre-configured credential resolver
#     that includes the default lookup chain for
#     credentials.
#     """
#     resolver = CredentialResolver(providers=[
#         BucketSnakeCredentialProvider(
#             load_config=lambda: session.full_config,
#             client_creator=session.create_client,
#             role_arn="arn:aws:iam::{account}:role/{app}-{account}".format(account=dest_account, app=source_app),
#             role_session_name=role_session_name,
#             source_credentials=source_provider
#         )
#     ])
#     return resolver


# def get_bucket_snake_botocore_session(source_app, dest_account, session_name, **botocore_session_kwargs):
#     session = Session(**botocore_session_kwargs)

#     # Get the current source credential provider:
#     source_provider = session._components.get_component("credential_provider")

#     # Overwrite the credential resolver:
#     session._components.lazy_register_component(
#         'credential_provider',
#         lambda: create_bucket_snake_credential_resolver(session, source_app, dest_account, session_name,
#                                                         source_provider))

#     return session


# class BucketSnakeAssumeRoleCredentialFetcher(AssumeRoleCredentialFetcher):
#     def __init__(self, client_creator, source_credentials_provider, role_arn, **kwargs):
#         self._source_credentials_provider = source_credentials_provider
#         super(BucketSnakeAssumeRoleCredentialFetcher, self).__init__(client_creator, self.credentials_provider,
#                                                                      role_arn, **kwargs)

#     @property
#     def credentials_provider(self):
#         return self._source_credentials_provider.load_credentials()


# class BucketSnakeCredentialProvider(AssumeRoleProvider):
#     METHOD = "bucket-snake"

#     def __init__(self, load_config, client_creator, role_arn, role_session_name, source_credentials):
#         self.cache = {}
#         self._load_config = load_config
#         self._client_creator = client_creator
#         self._loaded_config = {}
#         self.role_arn = role_arn
#         self.role_session_name = role_session_name
#         self.source_credentials = source_credentials    # This is a credential provider -- not source credentials!

#     def load(self):
#         self._loaded_config = self._load_config()
#         return self._load_creds_via_assume_role()

#     def _load_creds_via_assume_role(self):
#         extra_args = {"RoleSessionName": self.role_session_name}

#         fetcher = BucketSnakeAssumeRoleCredentialFetcher(
#             client_creator=self._client_creator,
#             source_credentials_provider=self.source_credentials,
#             role_arn=self.role_arn,
#             extra_args=extra_args,
#             cache=self.cache,
#         )
#         refresher = fetcher.fetch_credentials

#         # The initial credentials are empty and the expiration time is set
#         # to now so that we can delay the call to assume role until it is
#         # strictly needed.
#         return DeferredRefreshableCredentials(
#             method=self.METHOD,
#             refresh_using=refresher,
#             time_fetcher=_local_now
#         )


## TESTS:
# from botocore.session import Session

# from bucket_snake_client.credential_providers import get_bucket_snake_botocore_session, \
#     create_bucket_snake_credential_resolver


# def test_create_bucket_snake_credential_resolver():
#     session = Session()
#     source_provider = session._components.get_component("credential_provider")

#     resolver = create_bucket_snake_credential_resolver(session, "sourceApp", "123456789012", "bucketsnake", source_provider)

#     assert len(resolver.providers) == 1
#     assert resolver.providers[0].METHOD == "bucket-snake"
#     assert resolver.providers[0].role_arn == "arn:aws:iam::123456789012:role/sourceApp-123456789012"
#     assert resolver.providers[0].role_session_name == "bucketsnake"
#     assert resolver.providers[0].source_credentials == source_provider


# def test_get_bucket_snake_botocore_session(iam, sts):
#     session = get_bucket_snake_botocore_session("sourceApp", "123456789012", "bucketsnake")

#     creds = session.get_credentials()

#     assert creds.method == "bucket-snake"
#     assert creds.access_key
#     assert creds.secret_key
#     assert creds.token

test_credential_providers.py

from botocore.session import Session

from bucket_snake_client.credential_providers import get_bucket_snake_botocore_session, \
    create_bucket_snake_credential_resolver


def test_create_bucket_snake_credential_resolver():
    session = Session()
    source_provider = session._components.get_component("credential_provider")

    resolver = create_bucket_snake_credential_resolver(session, source_provider)

    assert len(resolver.providers) == 1
    assert resolver.providers[0].METHOD == "bucket-snake"
    assert resolver.providers[0]._credential_sourcer._providers == source_provider.providers


def test_get_bucket_snake_botocore_session(iam, sts):
    session = get_bucket_snake_botocore_session("sourceApp", "123456789012")

    creds = session.get_credentials()

    assert creds.method == "bucket-snake"
    assert creds.access_key
    assert creds.secret_key
    assert creds.token