import base64
import os

import boto3
from botocore.exceptions import ClientError
import xml.etree.ElementTree as ElementTree

from clokta.clokta_configuration import CloktaConfiguration
from clokta.common import Common
from clokta.awsrole import AwsRole


class AwsCredentialsGenerator:
    """
    This interfaces with AWS and, given a Okta-generated SAML token, can
    generate AWS credentials for a given role
    """

    def __init__(self, data_dir, clokta_config, saml_assertion):
        """
        Creates a credential generator capable of creating credentials in AWS
        :param clokta_config: the clokta configuration with which role has been designated default
        :type clokta_config: CloktaConfiguration
        :param saml_assertion: the saml token generated by Okta.  This will define what account and roles will be used
        :type saml_assertion: str
        """
        self.clokta_config = clokta_config
        self.saml_assertion = saml_assertion
        self.data_dir = data_dir
        self.bash_file = None  # type: str
        self.docker_file = None  # type: str
        self.roles = self.__deduce_roles_from_saml()  # type: [AwsRole]

        # We need to make sure, when interacting with AWS, we don't try to use
        # default creds as we are creating our own.
        envars_to_clear = [
            "AWS_PROFILE", 
            "AWS_ACCESS_KEY_ID",
            "AWS_SECRET_ACCESS_KEY",
            "AWS_SESSION_TOKEN"
        ]
        for envvar in envars_to_clear:
            if envvar in os.environ:
                del os.environ[envvar]


    def get_roles(self):
        """
        Returns the possible roles that can be assumed with the SAML token
        :return: list of roles
        :rtype: [AwsRole]
        """
        return self.roles

    def generate_creds(self, role):
        """
        :param role: the AWS role the user wants to assume
        :type role: AwsRole
        """

        client = boto3.client('sts')
        # Try for a 12 hour session.  If it fails, try for shorter periods
        assumed_role_credentials = None
        durations = [43200, 14400, 3600]
        for duration in durations:
            try:
                assumed_role_credentials = client.assume_role_with_saml(
                    RoleArn=role.role_arn,
                    PrincipalArn=role.idp_arn,
                    SAMLAssertion=self.saml_assertion,
                    DurationSeconds=duration
                )
                if duration == 3600:
                    Common.echo(message='YOUR SESSION WILL ONLY LAST ONE HOUR')
                break
            except ClientError as e:
                # If we get a validation error and we have shorter durations to try, try a shorter duration
                if e.response['Error']['Code'] != 'ValidationError' or duration == durations[-1]:
                    raise

        self.clokta_config.apply_credentials(credentials=assumed_role_credentials)
        self.bash_file = self.__write_sourceable_file(credentials=assumed_role_credentials)
        self.docker_file = self.__write_dockerenv_file(credentials=assumed_role_credentials)

    def __deduce_roles_from_saml(self):
        """
        Parse the SAML Token which contains the roles that can be assumed
        :return: roles contained in SAML token
        :rtype: [Role]
        """
        roles = []
        decoded_assertion = base64.b64decode(self.saml_assertion)
        root = ElementTree.fromstring(decoded_assertion)
        for saml2_attribute in root.iter('{urn:oasis:names:tc:SAML:2.0:assertion}Attribute'):
            if saml2_attribute.get('Name') == 'https://aws.amazon.com/SAML/Attributes/Role':
                for saml2_attr_value in saml2_attribute.iter('{urn:oasis:names:tc:SAML:2.0:assertion}AttributeValue'):
                    next_role = AwsRole(saml2_attr_value.text)
                    roles.append(next_role)
        return roles

    def __write_sourceable_file(self, credentials):
        """
        Generates a shell script to source in order to apply credentials to the shell environment.
        """
        creds = credentials['Credentials']
        output_file_name = '{dir}{profile}.sh'.format(
            dir=os.path.expanduser(self.data_dir),
            profile=self.clokta_config.profile_name
        )
        lines = [
            'export AWS_ACCESS_KEY_ID={}\n'.format(creds['AccessKeyId']),
            'export AWS_SECRET_ACCESS_KEY={}\n'.format(creds['SecretAccessKey'])
        ]
        if 'SessionToken' in creds:
            lines.append('export AWS_SESSION_TOKEN={}\n'.format(creds['SessionToken']))
        else:
            lines.append('unset AWS_SESSION_TOKEN')

        with open(output_file_name, mode='w') as file_handle:
            file_handle.writelines(lines)

        short_output_file_name = '{dir}{profile}.sh'.format(
            dir=self.data_dir,
            profile=self.clokta_config.profile_name
        )
        return short_output_file_name

    def __write_dockerenv_file(self, credentials):
        """
        Generates a Docker .env file that can be used with docker compose to inject into the environment.
        """
        creds = credentials['Credentials']
        output_file_name = '{dir}{profile}.env'.format(
            dir=os.path.expanduser(self.data_dir),
            profile=self.clokta_config.profile_name
        )
        lines = [
            'AWS_ACCESS_KEY_ID={}\n'.format(creds['AccessKeyId']),
            'AWS_SECRET_ACCESS_KEY={}\n'.format(creds['SecretAccessKey'])
        ]
        if 'SessionToken' in creds:
            lines.append('AWS_SESSION_TOKEN={}\n'.format(creds['SessionToken']))

        with open(output_file_name, mode='w') as file_handle:
            file_handle.writelines(lines)

        short_output_file_name = '{dir}{profile}.env'.format(
            dir=self.data_dir,
            profile=self.clokta_config.profile_name
        )
        return short_output_file_name
