#!/usr/bin/env python3

# Microsoft Azure Linux Agent
#
# Copyright 2018 Microsoft Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.
#

#
# This test adds extensions with multiple dependencies to a VMSS using the 'provisionAfterExtensions' property and
# validates they are enabled in order of dependencies.
#
import copy
import json
import random
import re
import uuid
from datetime import datetime
from typing import List, Dict, Any

from assertpy import fail
from azure.mgmt.compute.models import VirtualMachineScaleSetVMExtensionsSummary

from azurelinuxagent.common.future import UTC, datetime_min_utc

from tests_e2e.tests.ext_sequencing.ext_seq_test_cases import add_one_dependent_ext_without_settings, add_two_extensions_with_dependencies, \
    remove_one_dependent_extension, remove_all_dependencies, add_one_dependent_extension, \
    add_single_dependencies, remove_all_dependent_extensions, add_failing_dependent_extension_with_one_dependency, add_failing_dependent_extension_with_two_dependencies
from tests_e2e.tests.lib.agent_test import AgentVmssTest, TestSkipped
from tests_e2e.tests.lib.agent_test_context import AgentVmTestContext
from tests_e2e.tests.lib.virtual_machine_scale_set_client import VmssInstanceIpAddress
from tests_e2e.tests.lib.vm_extension_identifier import VmExtensionIds
from tests_e2e.tests.lib.logging import log
from tests_e2e.tests.lib.resource_group_client import ResourceGroupClient
from tests_e2e.tests.lib.ssh_client import SshClient


class ExtSequencing(AgentVmssTest):

    def __init__(self, context: AgentVmTestContext):
        super().__init__(context)
        self._scenario_start = datetime_min_utc

    # Cases to test different dependency scenarios
    _test_cases = [
        add_one_dependent_ext_without_settings,
        add_two_extensions_with_dependencies,
        # remove_one_dependent_extension should only be run after another test case which has RunCommandLinux in the
        # model
        remove_one_dependent_extension,
        # remove_all_dependencies should only be run after another test case which has extension dependencies in the
        # model
        remove_all_dependencies,
        add_one_dependent_extension,
        add_single_dependencies,
        # remove_all_dependent_extensions should only be run after another test case which has dependent extension in
        # the model
        remove_all_dependent_extensions,
        add_failing_dependent_extension_with_one_dependency,
        add_failing_dependent_extension_with_two_dependencies
    ]

    @staticmethod
    def _get_dependency_map(extensions: List[Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
        dependency_map: Dict[str, Dict[str, Any]] = {}

        for ext in extensions:
            ext_name = ext['name']
            provisioned_after = ext['properties'].get('provisionAfterExtensions')
            depends_on = provisioned_after if provisioned_after else []
            # We know an extension should fail if commandToExecute is exactly "exit 1"
            ext_settings = ext['properties'].get("settings")
            ext_command = ext['properties']['settings'].get("commandToExecute") if ext_settings else None
            should_fail = ext_command == "exit 1"
            dependency_map[ext_name] = {"should_fail": should_fail, "depends_on": depends_on}

        return dependency_map

    @staticmethod
    def _get_sorted_extension_names(extensions: List[VirtualMachineScaleSetVMExtensionsSummary], ssh_client: SshClient, test_case_start: datetime) -> List[str]:
        # Using VmExtensionIds to get publisher for each ext to be used in remote script
        extension_full_names = {
            "AzureMonitorLinuxAgent": VmExtensionIds.AzureMonitorLinuxAgent,
            "RunCommandLinux": VmExtensionIds.RunCommand,
            "CustomScript": VmExtensionIds.CustomScript
        }
        enabled_times = []
        for ext in extensions:
            # Only check extensions which succeeded provisioning
            if "succeeded" in ext.statuses_summary[0].code:
                enabled_time = ssh_client.run_command(f"ext_sequencing-get_ext_enable_time.py --ext '{extension_full_names[ext.name]}' --after_time '{test_case_start}'", use_sudo=True)
                formatted_time = datetime.strptime(enabled_time.strip(), u'%Y-%m-%dT%H:%M:%S.%fZ')
                enabled_times.append(
                    {
                        "name": ext.name,
                        "enabled_time": formatted_time
                     }
                )

        # sort the extensions based on their enabled datetime
        sorted_extensions = sorted(enabled_times, key=lambda ext_: ext_["enabled_time"])
        log.info("")
        log.info("Extensions sorted by time they were enabled: {0}".format(
            ', '.join(["{0}: {1}".format(ext["name"], ext["enabled_time"]) for ext in sorted_extensions])))
        sorted_extension_names = [ext["name"] for ext in sorted_extensions]
        return sorted_extension_names

    @staticmethod
    def _validate_extension_sequencing(dependency_map: Dict[str, Dict[str, Any]], sorted_extension_names: List[str], relax_check: bool):
        installed_ext = {}

        # Iterate through the extensions in the enabled order and validate if their depending extensions are already
        # enabled prior to that.
        for ext in sorted_extension_names:
            # Check if the depending extension are already installed
            if ext not in dependency_map:
                # There should not be any unexpected extensions on the scale set, even in the case we share the VMSS,
                # because we update the scale set model with the extensions. Any extensions that are not in the scale
                # set model would be disabled.
                fail("Unwanted extension found in VMSS Instance view: {0}".format(ext))
            if dependency_map[ext] is not None:
                dependencies = dependency_map[ext].get('depends_on')
                for dep in dependencies:
                    if installed_ext.get(dep) is None:
                        # The depending extension is not installed prior to the current extension
                        if relax_check:
                            log.info("{0} is not installed prior to {1}".format(dep, ext))
                        else:
                            fail("{0} is not installed prior to {1}".format(dep, ext))

            # Mark the current extension as installed
            installed_ext[ext] = ext

        # Validate that only extensions expected to fail, and their dependent extensions, failed
        for ext, details in dependency_map.items():
            failing_ext_dependencies = [dep for dep in details['depends_on'] if dependency_map[dep]['should_fail']]
            if ext not in installed_ext:
                if details['should_fail']:
                    log.info("Extension {0} failed as expected".format(ext))
                elif failing_ext_dependencies:
                    log.info("Extension {0} failed as expected because it is dependent on {1}".format(ext, ' and '.join(failing_ext_dependencies)))
                else:
                    fail("{0} unexpectedly failed. Only extensions that are expected to fail or depend on a failing extension should fail".format(ext))

        log.info("Validated extension sequencing")

    def run(self):
        instances_ip_address: List[VmssInstanceIpAddress] = self._context.vmss.get_instances_ip_address()
        ssh_clients: Dict[str, SshClient] = {}
        for instance in instances_ip_address:
            ssh_clients[instance.instance_name] = SshClient(ip_address=instance.ip_address, username=self._context.username, identity_file=self._context.identity_file)

        if not VmExtensionIds.AzureMonitorLinuxAgent.supports_distro(next(iter(ssh_clients.values())).run_command("get_distro.py").rstrip()):
            raise TestSkipped("Currently AzureMonitorLinuxAgent is not supported on this distro")

        # This is the base ARM template that's used for deploying extensions for this scenario
        base_extension_template = {
            "$schema": "https://schema.management.azure.com/schemas/2015-01-01/deploymentTemplate.json",
            "contentVersion": "1.0.0.0",
            "resources": [
                {
                    "type": "Microsoft.Compute/virtualMachineScaleSets",
                    "name": f"{self._context.vmss.name}",
                    "location": "[resourceGroup().location]",
                    "apiVersion": "2018-06-01",
                    "properties": {
                        "virtualMachineProfile": {
                            "extensionProfile": {
                                "extensions": []
                            }
                        }
                    }
                }
            ]
        }

        for case in self._test_cases:
            test_case_start = random.choice(list(ssh_clients.values())).run_command("date --utc '+%Y-%m-%d %T'").rstrip()
            if self._scenario_start == datetime_min_utc:
                self._scenario_start = test_case_start

            # Assign unique guid to forceUpdateTag for each extension to make sure they're always unique to force CRP
            # to generate a new sequence number each time
            test_guid = str(uuid.uuid4())
            extensions = case()
            for ext in extensions:
                ext["properties"].update({
                    "forceUpdateTag": test_guid
                })

            # We update the extension template here with extensions that are specific to the scenario that we want to
            # test out
            log.info("")
            log.info("Test case: {0}".format(case.__name__.replace('_', ' ')))
            log.info("Test case start time: {0}".format(test_case_start))
            ext_template = copy.deepcopy(base_extension_template)
            ext_template['resources'][0]['properties']['virtualMachineProfile']['extensionProfile'][
                'extensions'] = extensions

            # Log the dependency map for the extensions in this test case
            dependency_map = self._get_dependency_map(extensions)
            log.info("")
            log.info("The dependency map of the extensions for this test case is:")
            for ext, details in dependency_map.items():
                dependencies = details.get('depends_on')
                dependency_list = "-" if not dependencies else ' and '.join(dependencies)
                log.info("{0} depends on {1}".format(ext, dependency_list))

            # Deploy updated extension template to the scale set.
            log.info("")
            log.info("Deploying extensions with the above dependencies to the scale set...")
            rg_client = ResourceGroupClient(self._context.vmss.cloud, self._context.vmss.subscription,
                                            self._context.vmss.resource_group, self._context.vmss.location)
            try:
                rg_client.deploy_template(template=ext_template)
            except Exception as e:
                # We only expect to catch an exception during deployment if we are forcing one of the extensions to
                # fail. We know an extension should fail if "failing" is in the case name. Otherwise, report the
                # failure.
                #
                # Example deployment failure:
                #   (DeploymentFailed) At least one resource deployment operation failed. Please list deployment operations for details. Please see https://aka.ms/arm-deployment-operations for usage details.
                #   Code: DeploymentFailed
                #   Message: At least one resource deployment operation failed. Please list deployment operations for details. Please see https://aka.ms/arm-deployment-operations for usage details.
                #   Exception Details:	(Conflict) {
                # 	  "status": "Failed",
                # 	  "error": {
                # 	    "code": "ResourceDeploymentFailure",
                # 	    "message": "The resource write operation failed to complete successfully, because it reached terminal provisioning state 'Failed'.",
                # 	    "details": [
                # 	      {
                # 	        "code": "VMExtensionProvisioningError",
                # 	        "target": "0",
                # 	        "message": "Multiple VM extensions failed to be provisioned on the VM. The Extensions failed due to the errors: \n[Extension Name: 'AzureMonitorLinuxAgent'\nError Message: VM Extension 'AzureMonitorLinuxAgent' is marked as failed since it depends upon the VM Extension 'CustomScript' which has failed.]\n\n[Extension Name: 'CustomScript'\nError Message: VM has reported a failure when processing extension 'CustomScript' (publisher 'Microsoft.Azure.Extensions' and type 'CustomScript'). Error message: 'Enable failed: failed to execute command: command terminated with exit status=1\n[stdout]\n\n[stderr]\n'. More information on troubleshooting is available at https://aka.ms/VMExtensionCSELinuxTroubleshoot. ]\n"
                # 	      }
                # 	    ]
                # 	  }
                # 	}
                if "failing" not in case.__name__:
                    fail("Extension template deployment unexpectedly failed: {0}".format(e))
                else:
                    deployment_failure_pattern = r"[\s\S]*\"code\":\s*\"ResourceDeploymentFailure\"[\s\S]*\"details\":\s*\[\s*(?P<error>[\s\S]*)\]"
                    deployment_failure_match = re.match(deployment_failure_pattern, str(e))
                    try:
                        if deployment_failure_match is None:
                            raise Exception("Unable to match a ResourceDeploymentFailure")
                        error_json = json.loads(deployment_failure_match.group("error"))
                        error_code = error_json['code']
                        error_message = error_json['message']
                    except Exception as parse_exc:
                        fail("Extension template deployment failed as expected, but there was an error in parsing the failure. Parsing failure: {0}\nDeployment Failure: {1}".format(parse_exc, e))

                    msg_pattern = r"Multiple VM extensions failed to be provisioned on the VM[\s\S]*VM Extension '.*' is marked as failed since it depends upon the VM Extension 'CustomScript' which has failed."
                    if error_code != "VMExtensionProvisioningError" or re.match(msg_pattern, error_message) is None:
                        fail("Extension template deployment failed as expected, but with an unexpected error: {0}".format(e))

            # Get the extensions on the VMSS from the instance view
            log.info("")
            instance_view_extensions = self._context.vmss.get_instance_view().extensions

            # Validate that the extensions were enabled in the correct order on each instance of the scale set
            for instance_name, ssh_client in ssh_clients.items():
                log.info("")
                log.info("Validate extension sequencing on {0}:{1}...".format(instance_name, ssh_client.ip_address))

                # Sort the VM extensions by the time they were enabled
                sorted_extension_names = self._get_sorted_extension_names(instance_view_extensions, ssh_client, test_case_start)

                # Validate that the extensions were enabled in the correct order. We relax this check if no settings
                # are provided for a dependent extension, since the guest agent currently ignores dependencies in this
                # case.
                relax_check = True if "settings" in case.__name__ else False
                self._validate_extension_sequencing(dependency_map, sorted_extension_names, relax_check)

            log.info("------")

    def get_ignore_errors_before_timestamp(self) -> datetime:
        # Ignore errors in the agent log before the first test case starts
        if self._scenario_start == datetime_min_utc:
            return self._scenario_start
        return datetime.strptime(self._scenario_start, u'%Y-%m-%d %H:%M:%S').replace(tzinfo=UTC)

    def get_ignore_error_rules(self) -> List[Dict[str, Any]]:
        ignore_rules = [
            #
            # WARNING ExtHandler ExtHandler Missing dependsOnExtension on extension Microsoft.Azure.Monitor.AzureMonitorLinuxAgent
            # This message appears when an extension doesn't depend on another extension
            #
            {
                'message': r"Missing dependsOnExtension on extension .*"
            },
            #
            # WARNING ExtHandler ExtHandler Extension Microsoft.Azure.Monitor.AzureMonitorLinuxAgent does not have any settings. Will ignore dependency (dependency level: 1)
            # We currently ignore dependencies for extensions without settings
            #
            {
                'message': r"Extension .* does not have any settings\. Will ignore dependency \(dependency level: \d\)"
            },
            #
            # 2023-10-31T17:46:59.675959Z WARNING ExtHandler ExtHandler Dependent extension Microsoft.Azure.Extensions.CustomScript failed or timed out, will skip processing the rest of the extensions
            # We intentionally make CustomScript fail to test that dependent extensions are skipped
            #
            {
                'message': r"Dependent extension Microsoft.Azure.Extensions.CustomScript failed or timed out, will skip processing the rest of the extensions"
            },
            #
            # 2023-10-31T17:48:13.349214Z ERROR ExtHandler ExtHandler Event: name=Microsoft.Azure.Extensions.CustomScript, op=ExtensionProcessing, message=Dependent Extension Microsoft.Azure.Extensions.CustomScript did not succeed. Status was error, duration=0
            # We intentionally make CustomScript fail to test that dependent extensions are skipped
            #
            {
                'message': r"Event: name=Microsoft.Azure.Extensions.CustomScript, op=ExtensionProcessing, message=Dependent Extension Microsoft.Azure.Extensions.CustomScript did not succeed. Status was error, duration=0"
            },
            #
            # 2023-10-31T17:47:07.689083Z WARNING ExtHandler ExtHandler [PERIODIC] This status is being reported by the Guest Agent since no status file was reported by extension Microsoft.Azure.Monitor.AzureMonitorLinuxAgent: [ExtensionStatusError] Status file /var/lib/waagent/Microsoft.Azure.Monitor.AzureMonitorLinuxAgent-1.28.11/status/6.status does not exist
            # We expect extensions that are dependent on a failing extension to not report status
            #
            {
                'message': r"\[PERIODIC\] This status is being reported by the Guest Agent since no status file was reported by extension .*: \[ExtensionStatusError\] Status file \/var\/lib\/waagent\/.*\/status\/\d+.status does not exist"
            },
            #
            # 2023-10-31T17:48:11.306835Z WARNING ExtHandler ExtHandler A new goal state was received, but not all the extensions in the previous goal state have completed: [('Microsoft.Azure.Extensions.CustomScript', 'error'), ('Microsoft.Azure.Monitor.AzureMonitorLinuxAgent', 'transitioning'), ('Microsoft.CPlat.Core.RunCommandLinux', 'success')]
            # This message appears when the previous test scenario had failing extensions due to extension dependencies
            #
            {
                'message': r"A new goal state was received, but not all the extensions in the previous goal state have completed: \[(\(u?'.*', u?'(error|transitioning|success)'\),?)+\]"
            }
        ]
        return ignore_rules


if __name__ == "__main__":
    ExtSequencing.run_from_command_line()
