<commit_before>#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# authors:
# <NAME> <<EMAIL>>, <<EMAIL>>
# <NAME> <<EMAIL>>, <santiago[dot]cardenas[at]outlook[dot]com>
# <NAME> <<EMAIL>>,
# <NAME> <<EMAIL>>,
# <NAME> <<EMAIL>>
from __future__ import print_function

import boto3
import botocore
import logging
from threading import Lock
from time import sleep
from taskcat.exceptions import TaskCatException

log = logging.getLogger(__name__)


class ClientFactory(object):
    """Manages creating and caching boto3 clients, helpful when creating lots of
    clients in different regions or functions.

    Example usage:

    from tackcat import utils

    class MyClass(object):
        def __init__(self):
            self._boto_client = utils.ClientFactory()
        def my_function(self):
            s3_client = self._boto_client.get('s3', region='us-west-2')
            return s3_client.list_buckets()
    """

    def __init__(self, aws_access_key_id=None, aws_secret_access_key=None, aws_session_token=None, profile_name=None,
                 regional_cred_map=None):
        """Sets up the cache dict, a locking mechanism and the logging object

        Args:
            aws_access_key_id (str): [optional] IAM access key, defaults to None
            aws_secret_access_key (str): [optional] IAM secret key, defaults to None
            aws_session_token (str): [optional] IAM session token, defaults to None
            profile_name (str): [optional] credential profile to use, defaults to None
        """
        regional_cred_map = regional_cred_map if regional_cred_map else {}
        self._clients = {"default": {}}
        self._credential_sets = {}
        self._lock = Lock()
        self.put_credential_set('default', aws_access_key_id, aws_secret_access_key, aws_session_token, profile_name)
        for region_name, credential_dict in regional_cred_map.items():
            self.put_credential_set(region_name, **credential_dict)

        return

    def get_default_region(self, aws_access_key_id, aws_secret_access_key, aws_session_token, profile_name):
        """returns the default region for the credentials provided

        :param aws_access_key_id:
        :param aws_secret_access_key:
        :param aws_session_token:
        :param profile_name:
        :return:
        """
        region = 'us-east-1'
        try:
            if aws_access_key_id and aws_secret_access_key and aws_session_token:
                session = boto3.session.Session(aws_access_key_id=aws_access_key_id,
                                                aws_secret_access_key=aws_secret_access_key,
                                                aws_session_token=aws_session_token)
            elif aws_access_key_id and aws_secret_access_key:
                session = boto3.session.Session(aws_access_key_id=aws_access_key_id,
                                                aws_secret_access_key=aws_secret_access_key)
            elif profile_name:
                session = boto3.session.Session(profile_name=profile_name)
            else:
                session = boto3.session.Session()
            region = session.region_name
            if not region:
                log.warning("Region not set in credential chain, defaulting to us-east-1")
                region = 'us-east-1'
        except Exception as e:
            log.error('failed to get default region: %s' % str(e))
        finally:
            return region

    def put_credential_set(self, credential_set_name, aws_access_key_id=None, aws_secret_access_key=None,
                           aws_session_token=None, profile_name=None):
        """Adds or updates a credential set to be re-used when creating clients

                    aws_access_key_id (str): [optional] IAM access key, defaults to None
                    aws_secret_access_key (str): [optional] IAM secret key, defaults to None
                    aws_session_token (str): [optional] IAM session token, defaults to None
                    profile_name (str): [optional] credential profile to use, defaults to None
        """
        if (aws_access_key_id and not aws_secret_access_key) or (not aws_access_key_id and aws_secret_access_key):
            raise ValueError('"aws_access_key_id" and "aws_secret_access_key" must both be set')
        elif profile_name and (aws_access_key_id or aws_secret_access_key or aws_session_token):
            raise ValueError(
                '"profile_name" cannot be used with aws_access_key_id, aws_secret_access_key or aws_session_token')
        self._credential_sets[credential_set_name] = [aws_access_key_id, aws_secret_access_key, aws_session_token,
                                                      profile_name]

    def get(self, service, region=None, credential_set='default', aws_access_key_id=None,
            aws_secret_access_key=None, aws_session_token=None, s3v4=False, profile_name=None):
        """get a client for a given service and region, optionally with specific role, credentials and/or sig version

        Args:
            service (str): service name
            region (str): [optional] region name, defaults to current region
            credential_set (str): [optional] name used to seperate different sets of
                        credentials, defaults to "default" which uses either the auto-discovered
                        role, or the credentials configured when this class is instantiated
            aws_access_key_id (str): [optional] IAM access key, defaults to None
            aws_secret_access_key (str): [optional] IAM secret key, defaults to None
            aws_session_token (str): [optional] IAM session token, defaults to None
            s3v4 (bool): [optional] when True enables signature_version=s3v4 which is required for SSE
                         protected buckets/objects
            profile_name (str): [optional] credential profile to use, defaults to None
        Returns:
            class: boto3 client
        """
        if not aws_access_key_id and not profile_name:
            log.debug(
                "no explicit keys or profile for this client, fetching the credentials from the %s set" % credential_set
            )
            if credential_set not in self._credential_sets.keys():
                raise KeyError('credential set %s does not exist' % credential_set)
            if region and region in self._credential_sets.keys():
                credential_set = region
            aws_access_key_id, aws_secret_access_key, aws_session_token, profile_name = self._credential_sets[
                credential_set]
        if not region:
            region = self.get_default_region(aws_access_key_id, aws_secret_access_key, aws_session_token, profile_name)
        s3v4 = 's3v4' if s3v4 else 'default_sig_version'
        try:
            log.debug("Trying to get [%s][%s][%s][%s]" % (credential_set, region, service, s3v4))
            client = self._clients[credential_set][region][service][s3v4]
            if aws_access_key_id:
                if self._clients[credential_set][region]['session'].get_credentials().access_key != aws_access_key_id:
                    log.debug("credentials changed, forcing update...")
                    raise KeyError("New credentials for this credential_set, need a new session.")
            return client
        except KeyError:
            log.debug("Couldn't return an existing client, making a new one...")
            if credential_set not in self._clients.keys():
                self._clients[credential_set] = {}
            if region not in self._clients[credential_set].keys():
                self._clients[credential_set][region] = {}
            if service not in self._clients[credential_set].keys():
                self._clients[credential_set][region][service] = {}
            if 'session' not in self._clients[credential_set][region].keys():
                self._clients[credential_set][region]['session'] = self._create_session(
                    region, aws_access_key_id, aws_secret_access_key, aws_session_token, profile_name
                )
            self._clients[credential_set][region][service][s3v4] = self._create_client(
                credential_set, region, service, s3v4
            )
            return self._clients[credential_set][region][service][s3v4]

    def _create_session(self, region, access_key=None, secret_key=None, session_token=None, profile_name=None,
                        max_retries=4, delay=5, backoff_factor=2):
        """creates a boto3 session object

        Args:
            region (str): region name
            access_key (str): [optional] IAM secret key, defaults to None
            secret_key (str): [optional] IAM secret key, defaults to None
            session_token (str): [optional] IAM secret key, defaults to None
            profile_name (str): [optional] credential profile to use, defaults to None
            max_retries (int): [optional] number of retries, defaults to 4
            delay (int): [optional] retry delay in seconds, defaults to 5
            backoff_factor (int): [optional] retry delay exponent, defaults to 2
        """
        session = None
        retry = 0
        while not session:
            try:<commit_msg>Adding clientfactory_credset_name attribute to returned clients (ClientFactory)<commit_after>#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# authors:
# <NAME> <<EMAIL>>, <<EMAIL>>
# <NAME> <<EMAIL>>, <santiago[dot]cardenas[at]outlook[dot]com>
# <NAME> <<EMAIL>>,
# <NAME> <<EMAIL>>,
# <NAME> <<EMAIL>>
from __future__ import print_function

import boto3
import botocore
import logging
from threading import Lock
from time import sleep
from taskcat.exceptions import TaskCatException

log = logging.getLogger(__name__)


class ClientFactory(object):
    """Manages creating and caching boto3 clients, helpful when creating lots of
    clients in different regions or functions.

    Example usage:

    from tackcat import utils

    class MyClass(object):
        def __init__(self):
            self._boto_client = utils.ClientFactory()
        def my_function(self):
            s3_client = self._boto_client.get('s3', region='us-west-2')
            return s3_client.list_buckets()
    """

    def __init__(self, aws_access_key_id=None, aws_secret_access_key=None, aws_session_token=None, profile_name=None,
                 regional_cred_map=None):
        """Sets up the cache dict, a locking mechanism and the logging object

        Args:
            aws_access_key_id (str): [optional] IAM access key, defaults to None
            aws_secret_access_key (str): [optional] IAM secret key, defaults to None
            aws_session_token (str): [optional] IAM session token, defaults to None
            profile_name (str): [optional] credential profile to use, defaults to None
        """
        regional_cred_map = regional_cred_map if regional_cred_map else {}
        self._clients = {"default": {}}
        self._credential_sets = {}
        self._lock = Lock()
        self.put_credential_set('default', aws_access_key_id, aws_secret_access_key, aws_session_token, profile_name)
        for region_name, credential_dict in regional_cred_map.items():
            self.put_credential_set(region_name, **credential_dict)

        return

    def get_default_region(self, aws_access_key_id, aws_secret_access_key, aws_session_token, profile_name):
        """returns the default region for the credentials provided

        :param aws_access_key_id:
        :param aws_secret_access_key:
        :param aws_session_token:
        :param profile_name:
        :return:
        """
        region = 'us-east-1'
        try:
            if aws_access_key_id and aws_secret_access_key and aws_session_token:
                session = boto3.session.Session(aws_access_key_id=aws_access_key_id,
                                                aws_secret_access_key=aws_secret_access_key,
                                                aws_session_token=aws_session_token)
            elif aws_access_key_id and aws_secret_access_key:
                session = boto3.session.Session(aws_access_key_id=aws_access_key_id,
                                                aws_secret_access_key=aws_secret_access_key)
            elif profile_name:
                session = boto3.session.Session(profile_name=profile_name)
            else:
                session = boto3.session.Session()
            region = session.region_name
            if not region:
                log.warning("Region not set in credential chain, defaulting to us-east-1")
                region = 'us-east-1'
        except Exception as e:
            log.error('failed to get default region: %s' % str(e))
        finally:
            return region

    def put_credential_set(self, credential_set_name, aws_access_key_id=None, aws_secret_access_key=None,
                           aws_session_token=None, profile_name=None):
        """Adds or updates a credential set to be re-used when creating clients

                    aws_access_key_id (str): [optional] IAM access key, defaults to None
                    aws_secret_access_key (str): [optional] IAM secret key, defaults to None
                    aws_session_token (str): [optional] IAM session token, defaults to None
                    profile_name (str): [optional] credential profile to use, defaults to None
        """
        if (aws_access_key_id and not aws_secret_access_key) or (not aws_access_key_id and aws_secret_access_key):
            raise ValueError('"aws_access_key_id" and "aws_secret_access_key" must both be set')
        elif profile_name and (aws_access_key_id or aws_secret_access_key or aws_session_token):
            raise ValueError(
                '"profile_name" cannot be used with aws_access_key_id, aws_secret_access_key or aws_session_token')
        self._credential_sets[credential_set_name] = [aws_access_key_id, aws_secret_access_key, aws_session_token,
                                                      profile_name]

    def get(self, service, region=None, credential_set='default', aws_access_key_id=None,
            aws_secret_access_key=None, aws_session_token=None, s3v4=False, profile_name=None):
        """get a client for a given service and region, optionally with specific role, credentials and/or sig version

        Args:
            service (str): service name
            region (str): [optional] region name, defaults to current region
            credential_set (str): [optional] name used to seperate different sets of
                        credentials, defaults to "default" which uses either the auto-discovered
                        role, or the credentials configured when this class is instantiated
            aws_access_key_id (str): [optional] IAM access key, defaults to None
            aws_secret_access_key (str): [optional] IAM secret key, defaults to None
            aws_session_token (str): [optional] IAM session token, defaults to None
            s3v4 (bool): [optional] when True enables signature_version=s3v4 which is required for SSE
                         protected buckets/objects
            profile_name (str): [optional] credential profile to use, defaults to None
        Returns:
            class: boto3 client
        """
        if not aws_access_key_id and not profile_name:
            log.debug(
                "no explicit keys or profile for this client, fetching the credentials from the %s set" % credential_set
            )
            if credential_set not in self._credential_sets.keys():
                raise KeyError('credential set %s does not exist' % credential_set)
            if region and region in self._credential_sets.keys():
                credential_set = region
            aws_access_key_id, aws_secret_access_key, aws_session_token, profile_name = self._credential_sets[
                credential_set]
        if not region:
            region = self.get_default_region(aws_access_key_id, aws_secret_access_key, aws_session_token, profile_name)
        s3v4 = 's3v4' if s3v4 else 'default_sig_version'
        try:
            log.debug("Trying to get [%s][%s][%s][%s]" % (credential_set, region, service, s3v4))
            client = self._clients[credential_set][region][service][s3v4]
            if aws_access_key_id:
                if self._clients[credential_set][region]['session'].get_credentials().access_key != aws_access_key_id:
                    log.debug("credentials changed, forcing update...")
                    raise KeyError("New credentials for this credential_set, need a new session.")
        except KeyError:
            log.debug("Couldn't return an existing client, making a new one...")
            if credential_set not in self._clients.keys():
                self._clients[credential_set] = {}
            if region not in self._clients[credential_set].keys():
                self._clients[credential_set][region] = {}
            if service not in self._clients[credential_set].keys():
                self._clients[credential_set][region][service] = {}
            if 'session' not in self._clients[credential_set][region].keys():
                self._clients[credential_set][region]['session'] = self._create_session(
                    region, aws_access_key_id, aws_secret_access_key, aws_session_token, profile_name
                )
            self._clients[credential_set][region][service][s3v4] = self._create_client(
                credential_set, region, service, s3v4
            )
            client = self._clients[credential_set][region][service][s3v4]
        setattr(client, 'clientfactory_credset_name', credential_set)
        return client

    def _create_session(self, region, access_key=None, secret_key=None, session_token=None, profile_name=None,
                        max_retries=4, delay=5, backoff_factor=2):
        """creates a boto3 session object

        Args:
            region (str): region name
            access_key (str): [optional] IAM secret key, defaults to None
            secret_key (str): [optional] IAM secret key, defaults to None
            session_token (str): [optional] IAM secret key, defaults to None
            profile_name (str): [optional] credential profile to use, defaults to None
            max_retries (int): [optional] number of retries, defaults to 4
            delay (int): [optional] retry delay in seconds, defaults to 5
            backoff_factor (int): [optional] retry delay exponent, defaults to 2
        """
        session = None
        retry = 0
        while not session:
            try:
{'fraction_non_alphanumeric': 0.08007938210684402, 'fraction_numerical': 0.004827290281055567, 'mean_word_length': 2.4804928131416837, 'pattern_counts': {'":': 2, '<': 39, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4007385', 'n_tokens_mistral': 5161, 'n_tokens_neox': 5148, 'n_words': 1410}