Assume role for AWS cross-account with boto 3

boto3_assume_role.py

import logging
import boto3
import dateutil

logger = logging.getLogger()
logger.addHandler(logging.StreamHandler())  # Writes to console
logger.setLevel(logging.INFO)


def create_sts_client(aws_access_key_id=None,
                      aws_secret_access_key=None,
                      aws_session_token=None):
    # Create AWS security token service client
    sts_client = boto3.client(
        'sts',
        aws_access_key_id=aws_access_key_id,
        aws_secret_access_key=aws_secret_access_key,
        aws_session_token=aws_session_token)
    return sts_client


def get_crossaccount_credentials(access_key, secret_key, role_arn):
    # Create STS client and assume role with cross-account role
    client = create_sts_client(access_key, secret_key)
    return client.assume_role(RoleArn=role_arn, RoleSessionName='test-auth')


def get_ec2_client(access_key,
                   secret_key,
                   role_arn=None,
                   region_name='us-east-1'):
    aws_session_token = None
    if role_arn:
        # Getting temporary credentials AWS cross-account
        credentials = get_crossaccount_credentials(access_key, secret_key,
                                                   role_arn)
        access_key = credentials['Credentials']['AccessKeyId']
        secret_key = credentials['Credentials']['SecretAccessKey']
        aws_session_token = credentials['Credentials']['SessionToken']
        expiration = credentials['Credentials']['Expiration']
        expiration = expiration.astimezone(
            dateutil.tz.tzlocal()).strftime('%Y-%m-%d %H:%M:%S')
        logger.info("Retrieved creds from cross account. Valid till %s",
                    expiration)
    return boto3.client(
        'ec2',
        region_name=region_name,
        aws_access_key_id=access_key,
        aws_secret_access_key=secret_key,
        aws_session_token=aws_session_token)


if __name__ == '__main__':
    # Required values
    primary_aws_access_key_id = '<primary_account_access_key>'
    primary_aws_secret_access_key = '<primary_account_secret_key>'
    cross_account_role_arn = '<cross_account_role_arn>'

    # Creating EC2 client for cross account
    ec2_client = get_ec2_client(primary_aws_access_key_id,
                                primary_aws_secret_access_key,
                                cross_account_role_arn)
    logger.info(
        "Found %s instances",
        len(ec2_client.describe_instances()['Reservations'][0]['Instances']))