#!/usr/bin/env python

import unittest
import mock

from aws_saml_auth import amazon
from aws_saml_auth import configuration
from os import path
import os


class TestAmazon(unittest.TestCase):
    @property
    def valid_config(self):
        return configuration.Configuration(
            idp_id="IDPID",
            sp_id="SPID",
            username="user@example.com",
            password="hunter2",
        )

    def read_local_file(self, filename):
        here = path.abspath(path.dirname(__file__))
        with open(path.join(here, filename)) as fp:
            return fp.read().encode("utf-8")

    def test_sts_client(self):
        a = amazon.Amazon(self.valid_config, "dummy-encoded-saml")
        self.assertEqual(str(a.sts_client.__class__), "<class 'botocore.client.STS'>")

    def test_role_extraction(self):
        saml_xml = self.read_local_file("valid-response.xml")
        a = amazon.Amazon(self.valid_config, saml_xml)
        self.assertIsInstance(a.roles, dict)
        list_of_testing_roles = [
            "arn:aws:iam::123456789012:role/admin",
            "arn:aws:iam::123456789012:role/read-only",
            "arn:aws:iam::123456789012:role/test",
        ]
        self.assertEqual(sorted(list(a.roles.keys())), sorted(list_of_testing_roles))

    def test_role_extraction_too_many_commas(self):
        # See https://github.com/cevoaustralia/aws-google-auth/issues/12
        saml_xml = self.read_local_file("too-many-commas.xml")
        a = amazon.Amazon(self.valid_config, saml_xml)
        self.assertIsInstance(a.roles, dict)
        list_of_testing_roles = [
            "arn:aws:iam::123456789012:role/admin",
            "arn:aws:iam::123456789012:role/read-only",
            "arn:aws:iam::123456789012:role/test",
        ]
        self.assertEqual(sorted(list(a.roles.keys())), sorted(list_of_testing_roles))

    def test_invalid_saml_too_soon(self):
        saml_xml = self.read_local_file("saml-response-too-soon.xml")
        self.assertFalse(amazon.Amazon.is_valid_saml_assertion(saml_xml))

    def test_invalid_saml_too_late(self):
        saml_xml = self.read_local_file("saml-response-too-late.xml")
        self.assertFalse(amazon.Amazon.is_valid_saml_assertion(saml_xml))

    def test_invalid_saml_expired_before_valid(self):
        saml_xml = self.read_local_file("saml-response-expired-before-valid.xml")
        self.assertFalse(amazon.Amazon.is_valid_saml_assertion(saml_xml))

    def test_invalid_saml_bad_input(self):
        self.assertFalse(amazon.Amazon.is_valid_saml_assertion(None))
        self.assertFalse(amazon.Amazon.is_valid_saml_assertion("Malformed Base64"))
        self.assertFalse(amazon.Amazon.is_valid_saml_assertion(123456))
        self.assertFalse(amazon.Amazon.is_valid_saml_assertion(""))
        self.assertFalse(
            amazon.Amazon.is_valid_saml_assertion("QmFkIFhNTA==")
        )  # Bad XML

    def test_valid_saml(self):
        saml_xml = self.read_local_file("saml-response-no-expire.xml")
        self.assertTrue(amazon.Amazon.is_valid_saml_assertion(saml_xml))

    @mock.patch.dict(
        os.environ, {"AWS_PROFILE": "xxx-xxxx", "DEFAULT_AWS_PROFILE": "blart"}
    )
    def test_sts_client_with_invalid_profile(self):
        a = amazon.Amazon(self.valid_config, "dummy-encoded-saml")

        self.assertIsNotNone(a.sts_client)

        self.assertEqual("xxx-xxxx", os.environ["AWS_PROFILE"])
        self.assertEqual("blart", os.environ["DEFAULT_AWS_PROFILE"])
