def assume_role()

in enableDetective.py [0:0]


def assume_role(aws_account_number: str, role_name: str) -> boto3.Session:
    """
    Assumes the provided role in each account and returns a Detective client.

    Args:
        - aws_account_number: AWS Account Number
        - role_name: Role to assume in target account

    Returns:
        Detective client in the specified AWS Account and Region
    """
    try:
        # Beginning the assume role process for account
        sts_client = boto3.client('sts')

        # Get the current partition
        partition = sts_client.get_caller_identity()['Arn'].split(":")[1]

        response = sts_client.assume_role(
            RoleArn='arn:{}:iam::{}:role/{}'.format(
                partition,
                aws_account_number,
                role_name
            ),
            RoleSessionName='EnableDetective'
        )
        # Storing STS credentials
        session = boto3.Session(
            aws_access_key_id=response['Credentials']['AccessKeyId'],
            aws_secret_access_key=response['Credentials']['SecretAccessKey'],
            aws_session_token=response['Credentials']['SessionToken']
        )
    except Exception as e:
        logging.exception(f'exception: {e}')
        
    logging.info(f"Assumed session for {aws_account_number}.")

    return session