def assume_role()

in src/securityhub_enabler.py [0:0]

• 24 lines of code
• 2 McCabe index (conditional complexity)

def assume_role(aws_account_number, role_name):
    """
    Assumes the provided role in each account and returns a region_session object
    :param aws_account_number: AWS Account Number
    :param role_name: Role to assume in target account
    :return: Session object for the specified AWS Account and Region
    """
    sts_client = boto3.client(
        'sts',
        region_name=os.environ['AWS_REGION'],
        endpoint_url=f"https://sts.{os.environ['AWS_REGION']}.amazonaws.com"
        )
    partition = sts_client.get_caller_identity()['Arn'].split(":")[1]
    current_account = sts_client.get_caller_identity()['Arn'].split(":")[4]
    if aws_account_number == current_account:
        LOGGER.info(f"Using existing region_session for Account {aws_account_number}")
        return session
    else:
        response = sts_client.assume_role(
            RoleArn='arn:%s:iam::%s:role/%s' % (
                partition, aws_account_number, role_name),
            RoleSessionName='EnableSecurityHub'
        )
        sts_session = boto3.Session(
            aws_access_key_id=response['Credentials']['AccessKeyId'],
            aws_secret_access_key=response['Credentials']['SecretAccessKey'],
            aws_session_token=response['Credentials']['SessionToken']
        )
        LOGGER.info(f"Assumed region_session for Account {aws_account_number}")
        return sts_session