CreateAWSClientUsingSTSCredential.py

import boto3
import sys

# Please fill in the following variables
ROLE_ARN = ""
ROLE_SESSION_NAME = ""
EXTERNAL_ID = ""
DEFAULT_REGION = ""
 
 
class CreateAWSClientUsingSTSCredential(object):
    """
    CreateAWSClientUsingSTSCredential
 
    eg.
        c = CreateAWSClientUsingSTSCredential()
        c.get_sts_credentials()
        s3 = c.create_aws_client_using_sts_credentials('s3', 'resource')
    """
    def __init__(self):
        self.credentials = None
 
    def get_sts_credentials(self, role_arn=ROLE_ARN, role_session_name=ROLE_SESSION_NAME, external_id=EXTERNAL_ID):
        """
        type role_arn: 'string'
        type role_session_name: 'string'
        type external_id: 'string'
 
        optional:
            http://boto3.readthedocs.io/en/latest/reference/services/sts.html#STS.Client.assume_role
        """
        try:
            client = boto3.client('sts')
            self.credentials = client.assume_role(RoleArn=role_arn,
                                      RoleSessionName=role_session_name,
                                      ExternalId=external_id)['Credentials']
        except Exception as e:
            print(e)
 
 
    def create_aws_client_using_sts_credentials(self, aws_service, resource, region_name=DEFAULT_REGION):
        """
        type aws_service: string
        type resource: string
        """
        session = boto3.Session(aws_access_key_id=self.credentials['AccessKeyId'],
                                aws_secret_access_key=self.credentials['SecretAccessKey'],
                                aws_session_token=self.credentials['SessionToken'],
                                region_name=region_name)
        try:
            if resource == 'client':
                return session.client(aws_service)
            elif resource == 'resource':
                return session.resource(aws_service)
            else:
                raise NotImplementedError()
        except Exception as e:
            print(e)
            sys.exit(1)