boto3_session.py

'''
Description:
    - pivot into other accounts with specified RoleARN

References:
    assume_role:
        - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html#STS.Client.assume_role
'''
import boto3
import botocore.exceptions
from boto3.session import Session

def assume_role(pn, arn, session_name):
    '''
    Reference: assume_role
    Returns temporary security credentials used to access AWS resources
    that you might not normally have access to
    '''
    boto3.setup_default_session(profile_name=pn)
    client = boto3.client('sts')
    # account_info = client.get_caller_identity()
    # account_id = account_info["Account"]
    try:
        response = client.assume_role(RoleArn=arn, RoleSessionName=session_name, DurationSeconds=900)
    except botocore.exceptions.ClientError as error:
        print(error)
        return None
    session = Session(aws_access_key_id=response['Credentials']['AccessKeyId'], aws_secret_access_key=response['Credentials']['SecretAccessKey'], aws_session_token=response['Credentials']['SessionToken'])
    # sts_client = session.client('sts')
    # account_pivot = sts_client.get_caller_identity()
    return session


def boto3_session(SESSION, MODULE):
    '''
    Rate limiting:
    https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html
    '''
    config = Config(
       retries = {
          'max_attempts': 10,
          'mode': 'standard'
       }
    )
    return SESSION.client(MODULE, config=config)