# --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------

import os
import shutil
import tempfile
import unittest
import mock

from azure.cli.core.keys import is_valid_ssh_rsa_public_key
from azure.cli.command_modules.vm._validators import (validate_ssh_key,
                                                      _figure_out_storage_source,
                                                      _validate_admin_username,
                                                      _validate_admin_password,
                                                      _parse_image_argument,
                                                      process_disk_or_snapshot_create_namespace,
                                                      _validate_vmss_create_subnet,
                                                      _get_next_subnet_addr_suffix,
                                                      _validate_vm_vmss_msi,
                                                      _validate_vm_vmss_accelerated_networking)
from azure.cli.command_modules.vm._vm_utils import normalize_disk_info, update_disk_sku_info
from azure.cli.core.mock import DummyCli
from knack.util import CLIError


class TestActions(unittest.TestCase):
    def test_generate_specfied_ssh_key_files(self):
        temp_dir_name = tempfile.mkdtemp(prefix="ssh_dir_")

        # cleanup temporary directory and its contents
        self.addCleanup(shutil.rmtree, path=temp_dir_name)

        # first create file paths for the keys to be generated
        fd, private_key_file = tempfile.mkstemp(dir=temp_dir_name)
        os.close(fd)
        public_key_file = private_key_file + '.pub'
        os.remove(private_key_file)

        args = mock.MagicMock()
        args.ssh_key_value = [public_key_file]
        args.generate_ssh_keys = True

        # 1 verify we generate key files if not existing
        validate_ssh_key(args)

        generated_public_key_string = args.ssh_key_value[0]
        self.assertTrue(bool(args.ssh_key_value))
        self.assertTrue(is_valid_ssh_rsa_public_key(generated_public_key_string))
        self.assertTrue(os.path.isfile(private_key_file))

        # 2 verify we load existing key files
        # for convinience we will reuse the generated file in the previous step
        args2 = mock.MagicMock()
        args2.ssh_key_value = [generated_public_key_string]
        args2.generate_ssh_keys = False
        validate_ssh_key(args2)
        # we didn't regenerate
        self.assertEqual(generated_public_key_string, args.ssh_key_value[0])

        # 3 verify we do not generate unless told so
        fd, private_key_file2 = tempfile.mkstemp(dir=temp_dir_name)
        os.close(fd)
        public_key_file2 = private_key_file2 + '.pub'
        args3 = mock.MagicMock()
        args3.ssh_key_value = [public_key_file2]
        args3.generate_ssh_keys = False
        with self.assertRaises(CLIError):
            validate_ssh_key(args3)

        # 4 verify file naming if the pub file doesn't end with .pub
        fd, public_key_file4 = tempfile.mkstemp(dir=temp_dir_name)
        os.close(fd)
        public_key_file4 += '1'  # make it nonexisting
        args4 = mock.MagicMock()
        args4.ssh_key_value = [public_key_file4]
        args4.generate_ssh_keys = True
        validate_ssh_key(args4)
        self.assertTrue(os.path.isfile(public_key_file4 + '.private'))
        self.assertTrue(os.path.isfile(public_key_file4))

    def test_figure_out_storage_source(self):
        test_data = 'https://av123images.blob.core.windows.net/images/TDAZBET.vhd'
        src_blob_uri, src_disk, src_snapshot = _figure_out_storage_source(DummyCli(), 'tg1', test_data)
        self.assertFalse(src_disk)
        self.assertFalse(src_snapshot)
        self.assertEqual(src_blob_uri, test_data)

        test_data = '/subscriptions/0b1f6471-1bf0-4dda-aec3-cb9272f09590/resourceGroups/JAVACSMRG6017/providers/Microsoft.Compute/disks/ex.vhd'
        src_blob_uri, src_disk, src_snapshot = _figure_out_storage_source(None, 'tg1', test_data)
        self.assertEqual(src_disk, test_data)
        self.assertFalse(src_snapshot)
        self.assertFalse(src_blob_uri)

    def test_source_storage_account_err_case(self):
        np = mock.MagicMock()
        cmd = mock.MagicMock()
        cmd.cli_ctx = DummyCli()
        np._cmd = cmd
        np.source_storage_account_id = '/subscriptions/123/resourceGroups/ygsrc/providers/Microsoft.Storage/storageAccounts/s123'
        np.source = '/subscriptions/123/resourceGroups/yugangw/providers/Microsoft.Compute/disks/d2'

        # action (should throw)
        kwargs = {'namespace': np}
        with self.assertRaises(CLIError):
            process_disk_or_snapshot_create_namespace(cmd, **kwargs)

        # with blob uri, should be fine
        np.source = 'https://s1.blob.core.windows.net/vhds/s1.vhd'
        process_disk_or_snapshot_create_namespace(cmd, **kwargs)

    def test_validate_admin_username_linux(self):
        # pylint: disable=line-too-long
        err_invalid_char = r'admin user name cannot contain upper case character A-Z, special characters \/"[]:|<>+=;,?*@#()! or start with $ or -'

        self._verify_username_with_ex('!@#', 'linux', err_invalid_char)
        self._verify_username_with_ex('gue[', 'linux', err_invalid_char)
        self._verify_username_with_ex('Aguest', 'linux', err_invalid_char)
        self._verify_username_with_ex('-gguest', 'linux', err_invalid_char)
        self._verify_username_with_ex('', 'linux', 'admin user name can not be empty')
        self._verify_username_with_ex('guest', 'linux',
                                      "This user name 'guest' meets the general requirements, but is specifically disallowed for this image. Please try a different value.")

        _validate_admin_username('g-uest1', 'linux')
        _validate_admin_username('guest1', 'linux')
        _validate_admin_username('guest1.', 'linux')

    def test_validate_admin_username_windows(self):
        # pylint: disable=line-too-long
        err_invalid_char = r'admin user name cannot contain special characters \/"[]:|<>+=;,?*@# or ends with .'

        self._verify_username_with_ex('!@#', 'windows', err_invalid_char)
        self._verify_username_with_ex('gue[', 'windows', err_invalid_char)
        self._verify_username_with_ex('dddivid.', 'windows', err_invalid_char)
        self._verify_username_with_ex('backup', 'windows',
                                      "This user name 'backup' meets the general requirements, but is specifically disallowed for this image. Please try a different value.")

        _validate_admin_username('AGUEST', 'windows')
        _validate_admin_username('g-uest1', 'windows')
        _validate_admin_username('guest1', 'windows')

    def test_validate_admin_password_linux(self):
        # pylint: disable=line-too-long
        err_length = 'The password length must be between 12 and 72'
        err_variety = 'Password must have the 3 of the following: 1 lower case character, 1 upper case character, 1 number and 1 special character'

        self._verify_password_with_ex('te', 'linux', err_length)
        self._verify_password_with_ex('P12' + '3' * 70, 'linux', err_length)
        self._verify_password_with_ex('te12312312321', 'linux', err_variety)

        _validate_admin_password('Password22345', 'linux')
        _validate_admin_password('Password12!@#', 'linux')

    def test_validate_admin_password_windows(self):
        # pylint: disable=line-too-long
        err_length = 'The password length must be between 12 and 123'
        err_variety = 'Password must have the 3 of the following: 1 lower case character, 1 upper case character, 1 number and 1 special character'

        self._verify_password_with_ex('P1', 'windows', err_length)
        self._verify_password_with_ex('te14' + '3' * 120, 'windows', err_length)
        self._verify_password_with_ex('te12345678997', 'windows', err_variety)

        _validate_admin_password('Password22!!!', 'windows')
        _validate_admin_password('Pas' + '1' * 70, 'windows')

    def _verify_username_with_ex(self, admin_username, is_linux, expected_err):
        with self.assertRaises(CLIError) as context:
            _validate_admin_username(admin_username, is_linux)
        self.assertTrue(expected_err in str(context.exception))

    def _verify_password_with_ex(self, admin_password, is_linux, expected_err):
        with self.assertRaises(CLIError) as context:
            _validate_admin_password(admin_password, is_linux)
        self.assertTrue(expected_err in str(context.exception))

    @mock.patch('azure.cli.command_modules.vm._validators._compute_client_factory', autospec=True)
    def test_parse_image_argument(self, client_factory_mock):
        compute_client = mock.MagicMock()
        image = mock.MagicMock()
        cmd = mock.MagicMock()
        cmd.cli_ctx = DummyCli()
        image.plan.name = 'plan1'
        image.plan.product = 'product1'
        image.plan.publisher = 'publisher1'
        compute_client.virtual_machine_images.get.return_value = image
        client_factory_mock.return_value = compute_client

        np = mock.MagicMock()
        np.location = 'some region'
        np.plan_name, np.plan_publisher, np.plan_product = '', '', ''
        np.image = 'publisher1:offer1:sku1:1.0.0'

        # action
        _parse_image_argument(cmd, np)

        # assert
        self.assertEqual('plan1', np.plan_name)
        self.assertEqual('product1', np.plan_product)
        self.assertEqual('publisher1', np.plan_publisher)

    @mock.patch('azure.cli.command_modules.vm._validators._compute_client_factory', autospec=True)
    @mock.patch('azure.cli.command_modules.vm._validators.logger.warning', autospec=True)
    def test_parse_staging_image_argument(self, logger_mock, client_factory_mock):
        from msrestazure.azure_exceptions import CloudError
        compute_client = mock.MagicMock()
        resp = mock.MagicMock()
        cmd = mock.MagicMock()
        cmd.cli_ctx = DummyCli()
        resp.status_code = 404
        resp.text = '{"Message": "Not Found"}'

        compute_client.virtual_machine_images.get.side_effect = CloudError(resp, error='image not found')
        client_factory_mock.return_value = compute_client

        np = mock.MagicMock()
        np.location = 'some region'
        np.image = 'publisher1:offer1:sku1:1.0.0'
        np.plan_name, np.plan_publisher, np.plan_product = '', '', ''

        # action
        _parse_image_argument(cmd, np)

        # assert
        logger_mock.assert_called_with("Querying the image of '%s' failed for an error '%s'. "
                                       "Configuring plan settings will be skipped", 'publisher1:offer1:sku1:1.0.0',
                                       'image not found')

    def test_parse_unmanaged_image_argument(self):
        np = mock.MagicMock()
        np.image = 'https://foo.blob.core.windows.net/vhds/1'
        cmd = mock.MagicMock()
        # action & assert
        self.assertEqual(_parse_image_argument(cmd, np), 'uri')

    def test_parse_managed_image_argument(self):
        np = mock.MagicMock()
        np.image = '/subscriptions/123/resourceGroups/foo/providers/Microsoft.Compute/images/nixos-imag.vhd'
        cmd = mock.MagicMock()

        # action & assert
        self.assertEqual(_parse_image_argument(cmd, np), 'image_id')

    def test_get_next_subnet_addr_suffix(self):
        result = _get_next_subnet_addr_suffix('10.0.0.0/16', '10.0.0.0/24', 24)
        self.assertEqual(result, '10.0.1.0/24')

        # for 254~510 instances VMSS
        result = _get_next_subnet_addr_suffix('10.0.0.0/16', '10.0.0.0/23', 24)
        self.assertEqual(result, '10.0.2.0/24')

        # +1 overflows, so we go with -1
        result = _get_next_subnet_addr_suffix('12.0.0.0/16', '12.0.255.0/24', 24)
        self.assertEqual(result, '12.0.254.0/24')

        # handle carry bits to the next section
        result = _get_next_subnet_addr_suffix('12.0.0.0/15', '12.0.255.0/24', 24)
        self.assertEqual(result, '12.1.0.0/24')

        # error cases
        with self.assertRaises(CLIError):
            _get_next_subnet_addr_suffix('12.0.0.0/16', '12.0.255.0/15', 24)

        with self.assertRaises(CLIError):
            _get_next_subnet_addr_suffix('12.0.0.0/16', '12.1.0.0/16', 24)

        with self.assertRaises(CLIError):
            _get_next_subnet_addr_suffix('12.0.0.0/22', '12.0.0.0/22', 24)

        # verify end to end
        np_mock = mock.MagicMock()
        np_mock.vnet_type = 'new'
        np_mock.vnet_address_prefix = '10.0.0.0/16'
        np_mock.subnet_address_prefix = None
        np_mock.instance_count = 1000
        np_mock.app_gateway_type = 'new'
        np_mock.app_gateway_subnet_address_prefix = None
        np_mock.disable_overprovision = None
        _validate_vmss_create_subnet(np_mock)
        self.assertEqual(np_mock.app_gateway_subnet_address_prefix, '10.0.8.0/24')

    @mock.patch('azure.cli.command_modules.vm._validators._resolve_role_id', autospec=True)
    @mock.patch('azure.cli.core.commands.client_factory.get_subscription_id', autospec=True)
    def test_validate_msi_on_create(self, mock_get_subscription, mock_resolve_role_id):
        # check throw on : az vm/vmss create --assign-identity --role reader --scope ""
        np_mock = mock.MagicMock()
        cmd = mock.MagicMock()
        cmd.cli_ctx = DummyCli()
        np_mock.assign_identity = []
        np_mock.identity_scope = None
        np_mock.identity_role = 'reader'

        with self.assertRaises(CLIError) as err:
            _validate_vm_vmss_msi(cmd, np_mock)
        self.assertTrue("usage error: '--role reader' is not applicable as the '--scope' is "
                        "not provided" in str(err.exception))

        # check throw on : az vm/vmss create --scope "some scope"
        np_mock = mock.MagicMock()
        np_mock.assign_identity = None
        np_mock.identity_scope = 'foo-scope'
        with self.assertRaises(CLIError) as err:
            _validate_vm_vmss_msi(cmd, np_mock)
        self.assertTrue('usage error: --assign-identity [--scope SCOPE] [--role ROLE]' in str(err.exception))

        # check throw on : az vm/vmss create --role "reader"
        np_mock = mock.MagicMock()
        np_mock.assign_identity = None
        np_mock.identity_role = 'reader'
        with self.assertRaises(CLIError) as err:
            _validate_vm_vmss_msi(cmd, np_mock)
        self.assertTrue('usage error: --assign-identity [--scope SCOPE] [--role ROLE]' in str(err.exception))

        # check we set right role id
        np_mock = mock.MagicMock()
        np_mock.assign_identity = []
        np_mock.identity_scope = 'foo-scope'
        np_mock.identity_role = 'reader'
        mock_resolve_role_id.return_value = 'foo-role-id'
        _validate_vm_vmss_msi(cmd, np_mock)
        self.assertEqual(np_mock.identity_role_id, 'foo-role-id')
        self.assertEqual(np_mock.identity_role, 'reader')
        mock_resolve_role_id.assert_called_with(cmd.cli_ctx, 'reader', 'foo-scope')

    @mock.patch('azure.cli.command_modules.vm._validators._resolve_role_id', autospec=True)
    def test_validate_msi_on_assign_identity_command(self, mock_resolve_role_id):
        # check throw on : az vm/vmss assign-identity --role reader --scope ""
        np_mock = mock.MagicMock()
        cmd = mock.MagicMock()
        cmd.cli_ctx = DummyCli()
        np_mock.identity_scope = ''
        np_mock.identity_role = 'reader'

        with self.assertRaises(CLIError) as err:
            _validate_vm_vmss_msi(cmd, np_mock, from_set_command=True)
        self.assertTrue("usage error: '--role reader' is not applicable as the '--scope' is set to None",
                        str(err.exception))

        # check we set right role id
        np_mock = mock.MagicMock()
        np_mock.identity_scope = 'foo-scope'
        np_mock.identity_role = 'reader'
        np_mock.assign_identity = []
        mock_resolve_role_id.return_value = 'foo-role-id'
        _validate_vm_vmss_msi(cmd, np_mock, from_set_command=True)
        self.assertEqual(np_mock.identity_role_id, 'foo-role-id')
        mock_resolve_role_id.assert_called_with(cmd.cli_ctx, 'reader', 'foo-scope')

    def test_normalize_disk_info(self):
        from azure.cli.core.profiles._shared import AZURE_API_PROFILES, ResourceType
        from importlib import import_module

        api_version = AZURE_API_PROFILES['2019-03-01-hybrid'][ResourceType.MGMT_COMPUTE].default_api_version
        api_version = "v" + api_version.replace("-", "_")
        CachingTypes = import_module('azure.mgmt.compute.{}.models'.format(api_version)).CachingTypes

        cmd = mock.MagicMock()
        cmd.get_models.return_value = CachingTypes

        # verify caching configuring

        r = normalize_disk_info()
        self.assertEqual(r['os']['caching'], CachingTypes.read_write.value)

        r = normalize_disk_info(data_disk_cachings=['0=None'], data_disk_sizes_gb=[1, 2])
        self.assertEqual(r[0]['caching'], 'None')
        self.assertEqual(r[0]['diskSizeGB'], 1)
        self.assertEqual(r[1]['diskSizeGB'], 2)
        self.assertEqual(r['os']['caching'], CachingTypes.read_write.value)

        r = normalize_disk_info(data_disk_cachings=['None'], os_disk_caching='ReadOnly', data_disk_sizes_gb=[1, 2])
        self.assertEqual(r[0]['caching'], 'None')
        self.assertEqual(r[1]['caching'], 'None')
        self.assertEqual(r['os']['caching'], 'ReadOnly')

        r = normalize_disk_info(data_disk_cachings=['0=None', '1=ReadOnly'], data_disk_sizes_gb=[1, 2])
        self.assertEqual(r[0]['caching'], 'None')
        self.assertEqual(r[1]['caching'], 'ReadOnly')
        self.assertEqual(r['os']['caching'], CachingTypes.read_write.value)

        # CLI will not tweak any casing from the values
        r = normalize_disk_info(data_disk_cachings=['0=none', '1=readonly'], data_disk_sizes_gb=[1, 2])
        self.assertEqual(r[0]['caching'], 'none')
        self.assertEqual(r[1]['caching'], 'readonly')

        # error on configuring non-existing disks
        with self.assertRaises(CLIError) as err:
            normalize_disk_info(data_disk_cachings=['0=None', '1=foo'])
        self.assertTrue("Data disk with lun of '0' doesn't exist" in str(err.exception))

        # default to "None" across for Lv/Lv2 machines
        r = normalize_disk_info(data_disk_sizes_gb=[1], size='standard_L8s')
        self.assertEqual(r['os']['caching'], CachingTypes.none.value)
        self.assertEqual(r[0].get('caching'), None)

        # error on lv/lv2 machines with caching mode other than "None"
        with self.assertRaises(CLIError) as err:
            normalize_disk_info(data_disk_cachings=['ReadWrite'], data_disk_sizes_gb=[1, 2], size='standard_L16s_v2')
        self.assertTrue('for Lv series of machines, "None" is the only supported caching mode' in str(err.exception))

    @mock.patch('azure.cli.command_modules.vm._validators._compute_client_factory', autospec=True)
    def test_validate_vm_vmss_accelerated_networking(self, client_factory_mock):
        client_mock, size_mock = mock.MagicMock(), mock.MagicMock()
        client_mock.virtual_machine_sizes.list.return_value = [size_mock]
        client_factory_mock.return_value = client_mock
        # not a qualified size
        np = mock.MagicMock()
        np.size = 'Standard_Ds1_v2'
        np.accelerated_networking = None
        _validate_vm_vmss_accelerated_networking(None, np)
        self.assertIsNone(np.accelerated_networking)

        # qualified size and recognized distro
        np = mock.MagicMock()
        np.size = 'Standard_f8'
        size_mock.number_of_cores, size_mock.name = 8, 'Standard_f8'
        np.accelerated_networking = None
        np.os_publisher, np.os_offer, np.os_sku = 'Canonical', 'UbuntuServer', '16.04'
        _validate_vm_vmss_accelerated_networking(mock.MagicMock(), np)
        self.assertTrue(np.accelerated_networking)

        np = mock.MagicMock()
        np.size = 'Standard_DS4_v2'
        np.accelerated_networking = None
        np.os_publisher, np.os_offer, np.os_sku = 'coreos', 'coreos', 'alpha'
        size_mock.number_of_cores, size_mock.name = 8, 'Standard_DS4_v2'
        _validate_vm_vmss_accelerated_networking(mock.MagicMock(), np)
        self.assertTrue(np.accelerated_networking)

        np = mock.MagicMock()
        np.size = 'Standard_D3_v2'  # known supported 4 core size
        np.accelerated_networking = None
        np.os_publisher, np.os_offer, np.os_sku = 'coreos', 'coreos', 'alpha'
        _validate_vm_vmss_accelerated_networking(None, np)
        self.assertTrue(np.accelerated_networking)

        # not a qualified size, but user want it
        np = mock.MagicMock()
        np.size = 'Standard_Ds1_v2'
        np.accelerated_networking = True
        _validate_vm_vmss_accelerated_networking(None, np)
        self.assertTrue(np.accelerated_networking)

        # qualified size, but distro version not good
        np = mock.MagicMock()
        np.size = 'Standard_f8'
        size_mock.number_of_cores, size_mock.name = 8, 'Standard_f8'
        np.accelerated_networking = None
        np.os_publisher, np.os_offer, np.os_sku = 'canonical', 'UbuntuServer', '14.04.5-LTS'
        _validate_vm_vmss_accelerated_networking(mock.MagicMock(), np)
        self.assertIsNone(np.accelerated_networking)

        # qualified size, but distro infor is not available (say, custom images)
        np = mock.MagicMock()
        np.size = 'Standard_f8'
        size_mock.number_of_cores, size_mock.name = 8, 'Standard_f8'
        np.accelerated_networking = None
        np.os_publisher = None
        _validate_vm_vmss_accelerated_networking(mock.MagicMock(), np)
        self.assertIsNone(np.accelerated_networking)

        # qualified size, but distro version is not right
        np = mock.MagicMock()
        np.size = 'Standard_f8'
        size_mock.number_of_cores, size_mock.name = 8, 'Standard_f8'
        np.accelerated_networking = None
        np.os_publisher, np.os_offer, np.os_sku = 'oracle', 'oracle-linux', '7.3'
        _validate_vm_vmss_accelerated_networking(mock.MagicMock(), np)
        self.assertIsNone(np.accelerated_networking)

    def test_update_sku_from_dict(self):
        sku_tests = {"test_empty": ([""], {}),
                     "test_all": (["sku"], {"os": "sku", 1: "sku", 3: "sku"}),
                     "test_os": (["os=sku"], {"os": "sku"}),
                     "test_lun": (["1=sku"], {1: "sku"}),
                     "test_os_lun": (["os=sku", "1=sku_1"], {"os": "sku", 1: "sku_1"}),
                     "test_os_mult_lun": (["1=sku_1", "os=sku_os", "2=sku_2"],
                                          {1: "sku_1", "os": "sku_os", 2: "sku_2"}),
                     "test_double_equ": (["os==foo"], {"os": "=foo"}),
                     "test_err_no_eq": (["os=sku_1", "foo"], None),
                     "test_err_lone_eq": (["foo ="], None),
                     "test_err_float": (["2.7=foo"], None),
                     "test_err_bad_key": (["bad=foo"], None)}

        for test_sku, expected in sku_tests.values():
            if isinstance(expected, dict):
                # build info dict from expected values.
                info_dict = {lun: dict(managedDisk={'storageAccountType': None}) for lun in expected if lun != "os"}
                if "os" in expected:
                    info_dict["os"] = {}

                update_disk_sku_info(info_dict, test_sku)
                for lun in info_dict:
                    if lun == "os":
                        self.assertEqual(info_dict[lun]['storageAccountType'], expected[lun])
                    else:
                        self.assertEqual(info_dict[lun]['managedDisk']['storageAccountType'], expected[lun])
            elif expected is None:
                dummy_expected = ["os", 1, 2]
                info_dict = {lun: dict(managedDisk={'storageAccountType': None}) for lun in dummy_expected if lun != "os"}
                if "os" in dummy_expected:
                    info_dict["os"] = {}

                with self.assertRaises(CLIError):
                    update_disk_sku_info(info_dict, dummy_expected)
            else:
                self.fail("Test Expected value should be a dict or None, instead it is {}.".format(expected))


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