Boto3 Management Session with Refreshable Assume Role

boto3_session_management_and_assume_role.py

from os import getenv

from boto3 import Session
from util import make_refreshable_assume_role_session


def example():
    management_session = Session(aws_access_key_id=getenv('AWS_ACCESS_KEY_ID'),
                                 aws_secret_access_key=getenv('AWS_SECRET_ACCESS_KEY'))
    assume_role_params = dict(
        RoleArn=getenv('AWS_ROLE_ARN'),
        RoleSessionName='ExampleSessionName',
    )

    assumed_session = make_refreshable_assume_role_session(management_session, assume_role_params)

    print(assumed_session.client('sts').get_caller_identity())


if __name__ == '__main__':
    example()

util.py

from boto3 import Session
from botocore.credentials import CredentialProvider, CredentialResolver, Credentials, \
    RefreshableCredentials, create_assume_role_refresher
from botocore.session import get_session


def make_refreshable_assume_role_session(management_session: Session, assume_role_params: dict) -> Session:
    provider = ManagementSessionWithRefreshableAssumeRoleProvider(management_session, assume_role_params)

    resolver = CredentialResolver(providers=[provider])

    botocore_session = get_session()
    botocore_session.register_component('credential_provider', resolver)

    return Session(botocore_session=botocore_session, region_name=management_session.region_name)


class ManagementSessionWithRefreshableAssumeRoleProvider(CredentialProvider):
    METHOD = 'management-assume-role'

    def __init__(self, management_session: Session, assume_role_params: dict):
        assert isinstance(management_session, Session)
        self._management_session = management_session

        assert isinstance(assume_role_params, dict)
        self._assume_role_params = assume_role_params

        super().__init__()

    def load(self) -> Credentials:
        refresh = create_assume_role_refresher(self._management_session.client('sts'),
                                               self._assume_role_params)

        return RefreshableCredentials.create_from_metadata(
            metadata=refresh(),
            refresh_using=refresh,
            method=self.METHOD)