Assumed role session chaining (with credential refreshing) for boto3

aws_assume_role.py

# *** WARNING ***
# This gist is no longer maintained
# It has been replaced by aws-assume-role-lib in PyPI
# Documentation at https://github.com/benkehoe/aws-assume-role-lib
# It is still a single-file library, you can find the stable version here:
# https://raw.githubusercontent.com/benkehoe/aws-assume-role-lib/stable/aws_assume_role_lib/aws_assume_role_lib.py
# (link also available in the docs)

# Copyright 2020 Ben Kehoe
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Assumed role session chaining (with credential refreshing) for boto3.

This library provides the ability to create boto3 Sessions that use a given
role, assuming that role from a parent session, in a way such that the child
session can refresh its credentials when they expire by implicitly calling
AssumeRole again."""

import typing
import json

import boto3
import botocore

__version__ = "1.0.0"

__all__ = ["assume_role"]

def assume_role(session: boto3.Session, RoleArn: str, *,
        RoleSessionName: str=None,
        PolicyArns: typing.List[typing.Dict[str, str]]=None,
        Policy: typing.Union[str, typing.Dict]=None,
        DurationSeconds: int=None,
        Tags: typing.List[typing.Dict[str, str]]=None,
        TransitiveTagKeys:typing.List[str]=None,
        ExternalId: str=None,
        SerialNumber: str=None,
        TokenCode: str=None,
        validate: bool=True,
        additional_kwargs: typing.Dict=None) -> boto3.Session:
    """Produce a boto3 Session using the given role, assuming it from the input session

    Unlike creating a session with the credentials returned directly
    by sts.AssumeRole, the returned session will refresh the credentials
    automatically when they expire by calling AssumeRole again.

    By default, the parameters and session credentials are checked so that errors
    can be raised at this point, rather than more confusingly when the first call
    is made using the child session. This can be disabled by setting validate=False.

    The parent session is available on the child session
    in the property assume_role_parent_session.

    Additional arguments for AssumeRole, if any are added in the future,
    can be passed in additional_kwargs."""

    botocore_session = session._session

    if Policy is not None and not isinstance(Policy, str):
        Policy = json.dumps(Policy)

    extra_args = {}
    if additional_kwargs:
        extra_args.update(additional_kwargs)

    for var_name in [
            "RoleSessionName",
            "PolicyArns",
            "Policy",
            "DurationSeconds",
            "Tags",
            "TransitiveTagKeys",
            "ExternalId",
            "SerialNumber",
            "TokenCode"]:
        value = locals()[var_name]
        if value is not None:
            extra_args[var_name] = value

    if validate:
        validate_args = extra_args.copy()
        validate_args["RoleArn"] = RoleArn
        if "RoleSessionName" not in validate_args:
            # this gets generated later if it's not present
            validate_args["RoleSessionName"] = "ToBeGenerated"
        shape = session.client("sts")._service_model.shape_for("AssumeRoleRequest")
        botocore.validate.validate_parameters(validate_args, shape)

        credentials = botocore_session.get_credentials()
        if not credentials:
            raise botocore.exceptions.NoCredentialsError

    assume_role_provider = AssumeRoleProvider(
        botocore_session.create_client,
        botocore_session.get_credentials(),
        RoleArn,
        extra_args=extra_args,
    )

    assumed_role_botocore_session = botocore.session.Session()
    assumed_role_botocore_session.register_component(
        "credential_provider",
        botocore.credentials.CredentialResolver([assume_role_provider])
    )

    assumed_role_boto3_session = boto3.Session(botocore_session=assumed_role_botocore_session)

    assumed_role_boto3_session.assume_role_parent_session = session

    return assumed_role_boto3_session

class AssumeRoleProvider(botocore.credentials.CredentialProvider):
    METHOD = "assume-role"

    def __init__(self, client_creator, credential_loader, role_arn,
            extra_args=None, cache=None):
        self._client_creator = client_creator
        self._credential_loader = credential_loader
        self._role_arn = role_arn
        self._extra_args = extra_args
        self._fetcher = None
        if cache is None:
            cache = {}
        self._cache = cache

    def _get_fetcher(self):
        if not self._fetcher:
            self._fetcher = botocore.credentials.AssumeRoleCredentialFetcher(
                self._client_creator,
                self._credential_loader,
                self._role_arn,
                extra_args=self._extra_args,
                cache=self._cache
            )
        return self._fetcher

    def load(self):
        return botocore.credentials.DeferredRefreshableCredentials(
            self._get_fetcher().fetch_credentials,
            self.METHOD
        )

I have created a framework using boto3 in python for AWS migration. Same framework is being used by number of teams in my organization. Now the framework is working for different profile invocations.
But now currently we are implementing aws lakeformation process to register and grant permissions.
i have a scenario now. please refer below.

1. With the parent account I have to ingest the data --> with the current process it works fine. .~ .aws/credntials/
2. After the completion of step 1 from the parent account itself I have to assume the role and comple the lakeformation account. how can I achive this one using assuem role with the above code.
3. If I configure the assume role .~ .aws/config/ will it impact the current process step#1 or assume role only be used during the stsclient invocation.