# Copyright 2015 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 sys
from argparse import Namespace
from socket import timeout

from awscli.customizations.codedeploy.install import Install
from awscli.customizations.codedeploy.systems import (
    RHEL,
    System,
    Ubuntu,
    Windows,
)
from awscli.customizations.exceptions import (
    ConfigurationError,
    ParamValidationError,
)
from awscli.testutils import mock, unittest


class TestInstall(unittest.TestCase):
    def setUp(self):
        self.region = 'us-east-1'
        self.config_file = 'config-file'
        self.installer = 'install'
        self.bucket = f'aws-codedeploy-{self.region}'
        self.key = f'latest/{self.installer}'
        self.agent_installer = f's3://{self.bucket}/{self.key}'

        self.system_patcher = mock.patch('platform.system')
        self.system = self.system_patcher.start()
        self.system.return_value = 'Linux'

        self.linux_distribution_patcher = mock.patch(
            'awscli.compat.linux_distribution'
        )
        self.linux_distribution = self.linux_distribution_patcher.start()
        self.linux_distribution.return_value = ('Ubuntu', '', '')

        self.urlopen_patcher = mock.patch(
            'awscli.customizations.codedeploy.utils.urlopen'
        )
        self.urlopen = self.urlopen_patcher.start()
        self.urlopen.side_effect = timeout('Not EC2 instance')

        self.geteuid_patcher = mock.patch('os.geteuid', create=True)
        self.geteuid = self.geteuid_patcher.start()
        self.geteuid.return_value = 0

        self.isfile_patcher = mock.patch('os.path.isfile')
        self.isfile = self.isfile_patcher.start()
        self.isfile.return_value = False

        self.makedirs_patcher = mock.patch('os.makedirs')
        self.makedirs = self.makedirs_patcher.start()

        self.copyfile_patcher = mock.patch('shutil.copyfile')
        self.copyfile = self.copyfile_patcher.start()

        self.open_patcher = mock.patch(
            'awscli.customizations.codedeploy.systems.open',
            mock.mock_open(),
            create=True,
        )
        self.open = self.open_patcher.start()

        self.args = Namespace()
        self.args.override_config = False
        self.args.config_file = self.config_file
        self.args.agent_installer = None

        self.globals = Namespace()
        self.globals.region = self.region

        self.body = 'install-script'
        self.reader = mock.MagicMock()
        self.reader.read.return_value = self.body
        self.s3 = mock.MagicMock()
        self.s3.get_object.return_value = {'Body': self.reader}

        self.session = mock.MagicMock()
        self.session.create_client.return_value = self.s3
        self.install = Install(self.session)

    def tearDown(self):
        self.system_patcher.stop()
        self.linux_distribution_patcher.stop()
        self.urlopen_patcher.stop()
        self.geteuid_patcher.stop()
        self.isfile_patcher.stop()
        self.makedirs_patcher.stop()
        self.copyfile_patcher.stop()
        self.open_patcher.stop()

    def test_install_throws_on_invalid_region(self):
        self.globals.region = None
        self.session.get_config_variable.return_value = None
        error_msg = 'Region not specified.'
        with self.assertRaisesRegex(ConfigurationError, error_msg):
            self.install._run_main(self.args, self.globals)

    def test_install_throws_on_unsupported_system(self):
        self.system.return_value = 'Unsupported'
        with self.assertRaisesRegex(
            RuntimeError, System.UNSUPPORTED_SYSTEM_MSG
        ):
            self.install._run_main(self.args, self.globals)

    def test_install_throws_on_ec2_instance(self):
        self.urlopen.side_effect = None
        with self.assertRaisesRegex(
            RuntimeError, 'Amazon EC2 instances are not supported.'
        ):
            self.install._run_main(self.args, self.globals)
        self.assertIn('system', self.args)
        self.assertTrue(isinstance(self.args.system, Ubuntu))

    def test_install_throws_on_non_administrator(self):
        self.geteuid.return_value = 1
        with self.assertRaisesRegex(
            RuntimeError, 'You must run this command as sudo.'
        ):
            self.install._run_main(self.args, self.globals)

    def test_install_throws_on_no_override_config(self):
        self.isfile.return_value = True
        self.args.override_config = False
        with self.assertRaisesRegex(
            RuntimeError,
            'The on-premises instance configuration file already exists. '
            'Specify --override-config to update the existing on-premises '
            'instance configuration file.',
        ):
            self.install._run_main(self.args, self.globals)

    def test_install_throws_on_invalid_agent_installer(self):
        self.args.agent_installer = 'invalid-s3-location'
        with self.assertRaisesRegex(
            ParamValidationError,
            '--agent-installer must specify the Amazon S3 URL format as '
            's3://<bucket>/<key>.',
        ):
            self.install._run_main(self.args, self.globals)

    @mock.patch.object(Ubuntu, 'install')
    def test_install_with_agent_installer(self, install):
        self.args.agent_installer = self.agent_installer
        self.install._run_main(self.args, self.globals)
        self.assertIn('bucket', self.args)
        self.assertEqual(self.bucket, self.args.bucket)
        self.assertIn('key', self.args)
        self.assertEqual(self.key, self.args.key)
        self.assertIn('installer', self.args)
        self.assertEqual(self.installer, self.args.installer)
        install.assert_called_with(self.args)

    @mock.patch.object(Ubuntu, 'install')
    def test_install_for_ubuntu(self, install):
        self.system.return_value = 'Linux'
        self.linux_distribution.return_value = ('Ubuntu', '', '')
        self.install._run_main(self.args, self.globals)
        self.assertIn('bucket', self.args)
        self.assertEqual(self.bucket, self.args.bucket)
        self.assertIn('key', self.args)
        self.assertEqual('latest/install', self.args.key)
        self.assertIn('installer', self.args)
        self.assertEqual('install', self.args.installer)
        self.makedirs.assert_called_with('/etc/codedeploy-agent/conf')
        self.copyfile.assset_called_with(
            'codedeploy.onpremises.yml',
            '/etc/codedeploy-agent/conf/codedeploy.onpremises.yml',
        )
        install.assert_called_with(self.args)

    @mock.patch.object(Windows, 'install')
    @mock.patch.object(Windows, 'validate_administrator')
    def test_install_for_windows(self, validate_administrator, install):
        self.system.return_value = 'Windows'
        self.install._run_main(self.args, self.globals)
        self.assertIn('bucket', self.args)
        self.assertEqual(self.bucket, self.args.bucket)
        self.assertIn('key', self.args)
        self.assertEqual('latest/codedeploy-agent.msi', self.args.key)
        self.assertIn('installer', self.args)
        self.assertEqual('codedeploy-agent.msi', self.args.installer)
        self.makedirs.assert_called_with(r'C:\ProgramData\Amazon\CodeDeploy')
        self.copyfile.assset_called_with(
            'conf.onpremises.yml',
            r'C:\ProgramData\Amazon\CodeDeploy\conf.onpremises.yml',
        )
        validate_administrator.assert_called_with()
        install.assert_called_with(self.args)


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