<commit_before>from __future__ import print_function, absolute_import, unicode_literals, division

import json
import logging
import socket

from alppaca import NoCredentialsFoundException
from alppaca.compat import OrderedDict

from boto.sts import connect_to_region


class AssumedRoleCredentialsProvider(object):
    def __init__(self, credentials_provider, role_to_assume, aws_proxy_host=None, aws_proxy_port=None):
        self.credentials_provider = credentials_provider
        self.role_to_assume = role_to_assume
        self.aws_proxy_host = aws_proxy_host
        self.aws_proxy_port = aws_proxy_port
        self.logger = logging.getLogger(__name__)

    def get_credentials_for_all_roles(self):
        original_credentials = self.credentials_provider.get_credentials_for_all_roles()
        if not original_credentials:
            raise NoCredentialsFoundException("Got no credentials from: " + str(self.credentials_provider))

        for role in original_credentials:
            access_key, secret_key, token = self.parse_credentials_json(original_credentials[role])
            return self.get_credentials_for_assumed_role(access_key, secret_key, token)

    @staticmethod
    def parse_credentials_json(credentials_json):
        credentials_dict = json.loads(credentials_json)
        return credentials_dict['AccessKeyId'], credentials_dict['SecretAccessKey'], credentials_dict['Token']

    def get_credentials_for_assumed_role(self, access_key, secret_key, token):
        results = OrderedDict()

        self.logger.info("Connecting to AWS region eu-central-1 ...")
        try:
            conn = connect_to_region('eu-central-1',
                                     aws_access_key_id=access_key,
                                     aws_secret_access_key=secret_key,
                                     security_token=token,
                                     proxy=self.aws_proxy_host,
                                     proxy_port=self.aws_proxy_port
                                     )
            try:
                response = conn.assume_role(role_arn=self.role_to_assume, role_session_name=self.get_session_name())
                self.logger.info("Successfully got credentials for role: %s", self.role_to_assume)
            finally:<commit_msg>CLOUD-257 #15 Add proper debug logging for assume role<commit_after>from __future__ import print_function, absolute_import, unicode_literals, division

import json
import logging
import socket

from alppaca import NoCredentialsFoundException
from alppaca.compat import OrderedDict

from boto.sts import connect_to_region


class AssumedRoleCredentialsProvider(object):
    def __init__(self, credentials_provider, role_to_assume, aws_proxy_host=None, aws_proxy_port=None):
        self.credentials_provider = credentials_provider
        self.role_to_assume = role_to_assume
        self.aws_proxy_host = aws_proxy_host
        self.aws_proxy_port = aws_proxy_port
        self.logger = logging.getLogger(__name__)

    def get_credentials_for_all_roles(self):
        self.logger.debug("Getting credentials for all roles...")
        original_credentials = self.credentials_provider.get_credentials_for_all_roles()
        if not original_credentials:
            raise NoCredentialsFoundException("Got no credentials from: " + str(self.credentials_provider))

        for role in original_credentials:
            self.logger.debug("Got credentials for role '%s':", role)
            access_key, secret_key, token = self.parse_credentials_json(original_credentials[role])
            self.logger.debug("Access Key: %s", access_key)
            self.logger.debug("Secret Key (partial): %s...", secret_key[:10])
            self.logger.debug("Token (partial): %s...", token[:10])
            return self.get_credentials_for_assumed_role(access_key, secret_key, token)

    @staticmethod
    def parse_credentials_json(credentials_json):
        credentials_dict = json.loads(credentials_json)
        return credentials_dict['AccessKeyId'], credentials_dict['SecretAccessKey'], credentials_dict['Token']

    def get_credentials_for_assumed_role(self, access_key, secret_key, token):
        results = OrderedDict()
        region = "eu-central-1"
        self.logger.debug("Connecting to AWS region %s ...", region)
        try:
            conn = connect_to_region(region,
                                     aws_access_key_id=access_key,
                                     aws_secret_access_key=secret_key,
                                     security_token=token,
                                     proxy=self.aws_proxy_host,
                                     proxy_port=self.aws_proxy_port
                                     )
            try:
                response = conn.assume_role(role_arn=self.role_to_assume, role_session_name=self.get_session_name())
                self.logger.info("Successfully got credentials for role: %s", self.role_to_assume)
            finally:
{'fraction_non_alphanumeric': 0.07339081612191699, 'fraction_numerical': 0.00240625626629236, 'mean_word_length': 2.7057949479940566, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3090809', 'n_tokens_mistral': 1373, 'n_tokens_neox': 1351, 'n_words': 252}