# Copyright 2014 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.

import mock
import awscli.customizations.emr.emrutils as emrutils
import awscli.customizations.emr.createdefaultroles as createdefaultroles
from botocore.compat import json
from botocore.vendored import requests
from tests.unit.customizations.emr import EMRBaseAWSCommandParamsTest as \
    BaseAWSCommandParamsTest


EC2_ROLE_NAME = "EMR_EC2_DefaultRole"
EMR_ROLE_NAME = "EMR_DefaultRole"

EC2_ROLE_POLICY = {
    "Statement": [
        {
            "Action": [
                "cloudwatch:*",
                "dynamodb:*",
                "ec2:Describe*",
                "elasticmapreduce:Describe*",
                "rds:Describe*",
                "s3:*",
                "sdb:*",
                "sns:*",
                "sqs:*"
            ],
            "Effect": "Allow",
            "Resource": ["*"]
        }
    ]
}

CREATE_EC2_ROLE_RESULT = {
    "Role":  {
        "AssumeRolePolicyDocument": {
            "Version": "2008-10-17",
            "Statement": [
                {
                    "Action": "sts:AssumeRole",
                    "Sid": "",
                    "Effect": "Allow",
                    "Principal": {
                        "Service": "ec2.amazonaws.com"
                    }
                }
            ]
        },
        "RoleId": "AROAJG7O4RNNSRINMF6DI",
        "CreateDate": "2014-05-01T23:47:14.552Z",
        "RoleName": EC2_ROLE_NAME,
        "Path": "/",
        "Arn": "arn:aws:iam::176430881729:role/"+EC2_ROLE_NAME
    }
}

CONSTRUCTED_RESULT_OUTPUT = [
    {
        "Role": CREATE_EC2_ROLE_RESULT['Role'],
        "RolePolicy": EC2_ROLE_POLICY
    }
]

http_response = requests.models.Response()
http_response.status_code = 200

CN_EC2_ROLE_ARN = ('arn:aws-cn:iam::aws:policy/service-role/'
                   'AmazonElasticMapReduceforEC2Role')
US_GOV_EC2_ROLE_ARN = ('arn:aws-us-gov:iam::aws:policy/service-role/'
                       'AmazonElasticMapReduceforEC2Role')

EC2_ROLE_ARN = ('arn:aws:iam::aws:policy/service-role/'
                'AmazonElasticMapReduceforEC2Role')

CN_EMR_ROLE_ARN = ('arn:aws-cn:iam::aws:policy/service-role/'
                   'AmazonElasticMapReduceRole')

US_GOV_EMR_ROLE_ARN = ('arn:aws-us-gov:iam::aws:policy/'
                       'service-role/AmazonElasticMapReduceRole')

EMR_ROLE_ARN = ('arn:aws:iam::aws:policy/service-role/'
                'AmazonElasticMapReduceRole')


class TestCreateDefaultRole(BaseAWSCommandParamsTest):
    prefix = 'emr create-default-roles'

    ec2_role_policy_document = {
        "Version": "2008-10-17",
        "Statement": [
            {
                "Sid": "",
                "Effect": "Allow",
                "Principal": {"Service": "ec2.amazonaws.com.cn"},
                "Action": "sts:AssumeRole"
            }
        ]
    }

    emr_role_policy_document = {
        "Version": "2008-10-17",
        "Statement": [
            {
                "Sid": "",
                "Effect": "Allow",
                "Principal": {"Service": "elasticmapreduce.amazonaws.com.cn"},
                "Action": "sts:AssumeRole"
            }
        ]
    }

    def test_default_roles_exist(self):
        cmdline = self.prefix

        self.run_cmd(cmdline, expected_rc=0)
        self.assertEqual(len(self.operations_called), 3)

        self.assertEqual(self.operations_called[0][0].name, 'GetRole')
        self.assertEqual(self.operations_called[0][1]['RoleName'],
                         EC2_ROLE_NAME)

        self.assertEqual(self.operations_called[1][0].name,
                         'GetInstanceProfile')
        self.assertEqual(self.operations_called[1][1]['InstanceProfileName'],
                         EC2_ROLE_NAME)
        self.assertEqual(self.operations_called[2][0].name, 'GetRole')
        self.assertEqual(self.operations_called[2][1]['RoleName'],
                         EMR_ROLE_NAME)

    @mock.patch('awscli.customizations.emr.emr.'
                'CreateDefaultRoles._construct_result')
    @mock.patch('awscli.customizations.emr.emr.'
                'CreateDefaultRoles._check_if_instance_profile_exists')
    @mock.patch('awscli.customizations.emr.emr.'
                'CreateDefaultRoles._check_if_role_exists')
    @mock.patch('awscli.customizations.emr.emr.'
                'CreateDefaultRoles._get_role_policy')
    def test_default_roles_not_exist(self, get_rp_patch,
                                     role_exists_patch,
                                     instance_profile_exists_patch,
                                     construct_result_patch):
        get_rp_patch.return_value = False
        role_exists_patch.return_value = False
        instance_profile_exists_patch.return_value = False
        construct_result_patch.return_value = []

        cmdline = self.prefix + ' --region cn-north-1'

        self.run_cmd(cmdline, expected_rc=0)

        # Only 8 operations will be called as we are mocking
        # _check_if_role_exists and _check_if_instance_profile_exists methods.
        self.assertEqual(len(self.operations_called), 6)

        self.assertEqual(self.operations_called[0][0].name, 'CreateRole')
        self.assertEqual(self.operations_called[0][1]['RoleName'],
                         EC2_ROLE_NAME)
        self.assertEqual(
            self.operations_called[0][1]['AssumeRolePolicyDocument'],
            emrutils.dict_to_string(self.ec2_role_policy_document))

        self.assertEqual(self.operations_called[1][0].name,
                         'AttachRolePolicy')
        self.assertEqual(self.operations_called[1][1]['PolicyArn'],
                         CN_EC2_ROLE_ARN)
        self.assertEqual(self.operations_called[1][1]['RoleName'],
                         EC2_ROLE_NAME)

        self.assertEqual(self.operations_called[2][0].name,
                         'CreateInstanceProfile')
        self.assertEqual(self.operations_called[2][1]['InstanceProfileName'],
                         EC2_ROLE_NAME)

        self.assertEqual(self.operations_called[3][0].name,
                         'AddRoleToInstanceProfile')
        self.assertEqual(self.operations_called[3][1]['InstanceProfileName'],
                         EC2_ROLE_NAME)
        self.assertEqual(self.operations_called[3][1]['RoleName'],
                         EC2_ROLE_NAME)

        self.assertEqual(self.operations_called[4][0].name, 'CreateRole')
        self.assertEqual(self.operations_called[4][1]['RoleName'],
                         EMR_ROLE_NAME)
        self.assertEqual(
            self.operations_called[4][1]['AssumeRolePolicyDocument'],
            emrutils.dict_to_string(self.emr_role_policy_document))

        self.assertEqual(self.operations_called[5][0].name, 'AttachRolePolicy')
        self.assertEqual(self.operations_called[5][1]['PolicyArn'],
                         CN_EMR_ROLE_ARN)
        self.assertEqual(self.operations_called[5][1]['RoleName'],
                         EMR_ROLE_NAME)

    @mock.patch('awscli.customizations.emr.emr.'
                'CreateDefaultRoles._construct_result')
    @mock.patch('awscli.customizations.emr.createdefaultroles'
                '.get_service_principal')
    @mock.patch('awscli.customizations.emr.emr.'
                'CreateDefaultRoles._check_if_instance_profile_exists')
    @mock.patch('awscli.customizations.emr.emr.'
                'CreateDefaultRoles._check_if_role_exists')
    @mock.patch('awscli.customizations.emr.emr.'
                'CreateDefaultRoles._get_role_policy')
    def test_get_service_principal_parameters(self, get_rp_patch,
                                              role_exists_patch,
                                              instance_profile_exists_patch,
                                              get_sp_patch,
                                              construct_result_patch):
        get_rp_patch.return_value = "blah"
        get_sp_patch.return_value = 'elasticmapreduce.amazonaws.abc'
        role_exists_patch.return_value = False
        instance_profile_exists_patch.return_value = False
        construct_result_patch.return_value = []

        endpoint_url = 'https://elasticmapreduce.abc'
        cmdline = self.prefix + ' --endpoint ' + endpoint_url
        self.run_cmd(cmdline, expected_rc=0)
        self.assertEquals(get_sp_patch.call_args[0][1], endpoint_url)

    @mock.patch('botocore.session.Session.create_client')
    def test_call_parameters(self, call_patch):
        cmdline = self.prefix + ' --region eu-west-1' + ' --no-verify-ssl'
        self.run_cmd(cmdline, expected_rc=0)
        self.assertEquals(call_patch.call_args[1]['region_name'], 'eu-west-1')
        self.assertEquals(call_patch.call_args[1]['verify'], False)

    @mock.patch('botocore.session.Session.create_client')
    def test_call_parameters_only_endpoint(self, call_patch):
        endpoint_arg = 'https://elasticmapreduce.us-unknown-1.amazonaws.com'
        cmdline = self.prefix + ' --endpoint ' + endpoint_arg
        self.run_cmd(cmdline, expected_rc=0)
        self.assertEquals(call_patch.call_args[1]['endpoint_url'], None)

    @mock.patch('botocore.session.Session.create_client')
    def test_call_parameters_only_iam_endpoint(self, call_patch):
        endpoint_arg = 'https://elasticmapreduce.us-unknown-1.amazonaws.com'
        cmdline = self.prefix + ' --iam-endpoint ' + endpoint_arg
        self.run_cmd(cmdline, expected_rc=0)
        self.assertEquals(call_patch.call_args[1]['endpoint_url'],
                          endpoint_arg)

    @mock.patch('awscli.customizations.emr.emr.'
                'CreateDefaultRoles._get_role_policy')
    @mock.patch('awscli.customizations.emr.emr.'
                'CreateDefaultRoles._create_role_with_role_policy')
    @mock.patch('awscli.customizations.emr.emr.'
                'CreateDefaultRoles._check_if_instance_profile_exists')
    @mock.patch('awscli.customizations.emr.emr.'
                'CreateDefaultRoles._check_if_role_exists')
    def test_constructed_result(self, role_exists_patch,
                                instance_profile_exists_patch,
                                create_role_patch,
                                get_role_policy_patch):
        role_exists_patch.side_effect = side_effect_of_check_if_role_exists
        instance_profile_exists_patch.return_value = False
        create_role_patch.return_value = CREATE_EC2_ROLE_RESULT
        get_role_policy_patch.return_value = EC2_ROLE_POLICY

        cmdline = self.prefix + ' --region cn-north-1'
        result = self.run_cmd(cmdline, 0)
        expected_output = json.dumps(CONSTRUCTED_RESULT_OUTPUT, indent=4) +\
            '\n'
        self.assertEquals(result[0], expected_output)

    def test_policy_arn_construction(self):
        self.assertEquals(
            createdefaultroles.get_ec2_role_policy_arn("cn-north-1"),
            CN_EC2_ROLE_ARN)
        self.assertEquals(
            createdefaultroles.get_ec2_role_policy_arn("us-gov-west-1"),
            US_GOV_EC2_ROLE_ARN)
        self.assertEquals(
            createdefaultroles.get_ec2_role_policy_arn("eu-west-1"),
            EC2_ROLE_ARN)
        self.assertEquals(
            createdefaultroles.get_service_role_policy_arn("cn-north-1"),
            CN_EMR_ROLE_ARN)
        self.assertEquals(
            createdefaultroles.get_service_role_policy_arn("us-gov-west-1"),
            US_GOV_EMR_ROLE_ARN)
        self.assertEquals(
            createdefaultroles.get_service_role_policy_arn("eu-west-1"),
            EMR_ROLE_ARN)


def side_effect_of_check_if_role_exists(*args, **kwargs):
    if args[0] == EC2_ROLE_NAME:
        return False
    else:
        return True


if __name__ == "__main__":
    unittest.main()
