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

# pylint: disable=line-too-long, too-many-locals, too-many-statements, broad-except, too-many-branches
import json
import re
import shlex
import timeit
import traceback
import requests

from knack.log import get_logger

from azure.cli.command_modules.vm.custom import get_vm, _is_linux_os
from azure.cli.command_modules.storage.storage_url_helpers import StorageResourceIdentifier
from azure.cli.core.azclierror import InvalidArgumentValueError
from azure.mgmt.core.tools import parse_resource_id
from .exceptions import AzCommandError, SkuNotAvailableError, UnmanagedDiskCopyError, WindowsOsNotAvailableError, RunScriptNotFoundForIdError, SkuDoesNotSupportHyperV, ScriptReturnsError, SupportingResourceNotFoundError, CommandCanceledByUserError

from .command_helper_class import command_helper
from .repair_utils import (
    _uses_managed_disk,
    _call_az_command,
    _clean_up_resources,
    _fetch_compatible_sku,
    _fetch_source_disk_controller_type,
    _fetch_sku_disk_controller_types,
    _select_repair_disk_controller_type,
    _list_resource_ids_in_rg,
    _get_repair_resource_tag,
    _validate_tags_for_command,
    _fetch_compatible_windows_os_urn,
    _fetch_matching_windows_os_urn,
    _fetch_run_script_map,
    _fetch_run_script_path,
    _process_ps_parameters,
    _process_bash_parameters,
    _parse_run_script_raw_logs,
    _check_script_succeeded,
    _fetch_disk_info,
    _invoke_run_command,
    _get_cloud_init_script,
    _select_distro_linux,
    _check_linux_hyperV_gen,
    _select_distro_linux_gen2,
    _set_repair_map_url,
    REPAIR_LIBRARY_FORK,
    REPAIR_LIBRARY_BRANCH,
    _is_gen2,
    _unlock_encrypted_vm_run,
    _create_repair_vm,
    _check_n_start_vm,
    _check_existing_rg,
    _fetch_architecture,
    _select_distro_linux_Arm64,
    _fetch_vm_security_profile_parameters,
    _fetch_osdisk_security_profile_parameters,
    _make_public_ip_name
)

logger = get_logger(__name__)

PREVIEW_URL_ERROR = ("Invalid preview url. Write full URL of map.json file. "
                     "example https://github.com/{user}/repair-script-library/blob/main/map.json. "
                     "The branch name must be a single path segment.")

# The driver downloads from https://github.com/<fork>/repair-script-library/tarball/<branch>/,
# so the repository name and a single-segment branch are both part of the contract.
PREVIEW_URL_PATTERN = re.compile(
    r'^https://github\.com/(?P<fork>[^/]+)/repair-script-library/(?:blob|tree)/(?P<branch>[^/]+)/map\.json$')


def _parse_preview_url(preview):
    """Extract the fork and branch from a preview map.json URL.

    The URL is read positionally, so a branch name containing a slash shifts the fork to the
    repository name and resolves to an entirely different GitHub organization. Reject anything
    that does not match the documented shape rather than downloading scripts from a repository
    the caller never named.
    """
    match = PREVIEW_URL_PATTERN.match(str(preview).strip())
    if not match:
        raise InvalidArgumentValueError(PREVIEW_URL_ERROR)
    return match.group('fork'), match.group('branch')


def _build_repo_params(preview, is_linux):
    """Build the run-command parameters telling the run driver which script library to download.

    The Linux driver receives parameters positionally, so the fork and branch are always sent for
    Linux. Omitting them there would both shift the positions of the script parameters that follow
    and leave the driver downloading the default library while the run id resolved from a fork.
    """
    fork_name = REPAIR_LIBRARY_FORK
    branch_name = REPAIR_LIBRARY_BRANCH

    if preview:
        fork_name, branch_name = _parse_preview_url(preview)
    elif not is_linux:
        # The Windows driver declares these as named parameters with the same defaults.
        return []

    return ['repo_fork="{}"'.format(fork_name), 'repo_branch="{}"'.format(branch_name)]


def _set_source_resource_context(command, source_vm, source_vm_instance_view=None, disk_controller_type=None):
    """Populate telemetry-only resource shape without allowing enrichment to fail a command."""
    source_controller = disk_controller_type
    if source_controller is None:
        try:
            # Shared helper falls back to an ARM query when the SDK does not model the field.
            source_controller = _fetch_source_disk_controller_type(source_vm)
        except Exception as exception:
            logger.debug('Could not determine source VM disk controller type for telemetry: %s', exception)

    hyperv_generation = None
    if source_vm_instance_view:
        try:
            hyperv_generation = 'V{}'.format(_is_gen2(source_vm_instance_view))
        except (AttributeError, TypeError, ValueError) as exception:
            logger.debug('Could not determine source VM Hyper-V generation for telemetry: %s', exception)

    if hasattr(command, 'set_resource_context'):
        hardware_profile = getattr(source_vm, 'hardware_profile', None)
        command.set_resource_context(
            os_family='linux' if _is_linux_os(source_vm) else 'windows',
            vm_size=getattr(hardware_profile, 'vm_size', None),
            disk_controller_type=source_controller,
            hyperv_generation=hyperv_generation)
    return source_controller


def _enrich_source_resource_context(command, cmd, source_vm, resource_group_name, vm_name):
    """Add instance-view telemetry when available without changing command success or failure."""
    try:
        source_vm_instance_view = get_vm(cmd, resource_group_name, vm_name, 'instanceView')
    except Exception as exception:
        logger.debug('Could not fetch source VM instance view for telemetry: %s', exception)
        source_vm_instance_view = None
    _set_source_resource_context(command, source_vm, source_vm_instance_view)


def _enrich_repair_controller_context(command, cmd, resource_group_name, vm_name):
    """Add the actual repair VM controller when available without changing command behavior."""
    if not hasattr(command, 'set_resource_context'):
        return
    try:
        repair_vm = get_vm(cmd, resource_group_name, vm_name)
        repair_controller = _fetch_source_disk_controller_type(repair_vm)
        command.set_resource_context(repair_vm_disk_controller_type=repair_controller)
    except Exception as exception:
        logger.debug('Could not fetch repair VM disk controller type for telemetry: %s', exception)


def create(cmd, vm_name, resource_group_name, repair_password=None, repair_username=None, repair_vm_name=None, copy_disk_name=None, repair_group_name=None, unlock_encrypted_vm=False, enable_nested=False, associate_public_ip=False, distro='ubuntu', encrypt_recovery_key="", disable_trusted_launch=False, os_disk_type=None, tags=None, copy_tags=False, size=None, disk_controller_type=None, yes=False):
    """
    This function creates a repair VM.

    Parameters:
    - cmd: The command to be executed.
    - vm_name: The name of the virtual machine.
    - resource_group_name: The name of the resource group.
    - repair_password: The password for the repair VM. If not provided, a default will be used.
    - repair_username: The username for the repair VM. If not provided, a default will be used.
    - repair_vm_name: The name of the repair VM. If not provided, a default will be used.
    - copy_disk_name: The name of the disk to be copied. If not provided, the OS disk of the source VM will be used.
    - repair_group_name: The name of the repair group. If not provided, a default will be used.
    - unlock_encrypted_vm: If True, the encrypted VM will be unlocked. Default is False.
    - enable_nested: If True, nested virtualization will be enabled. Default is False.
    - associate_public_ip: If True, a public IP will be associated with the VM. Default is False.
    - distro: The Linux distribution to use for the repair VM. Default is 'ubuntu'.
    - encrypt_recovery_key: The Bitlocker recovery key to use for encrypting the VM. Default is an empty string.
    - disable_trusted_launch: A flag parameter that, when used, sets the security type of the repair VM to Standard.
    - os_disk_type: Set the OS disk storage account type of the repair VM to the specified type. The default is PremiumSSD_LRS.
    - tags: Tags to apply to the repair VM. Should be a dictionary or a string in key[=value] format.
    - copy_tags: If True, tags will be copied from the source VM to the repair VM. Default is False.
    - size: The size of the repair VM. If not provided, the size of the broken vm will be used.
    - disk_controller_type: Optional SCSI or NVMe override for the repair VM disk controller.
    """

    # Logging all the command parameters, except the sensitive data.
    # Mask sensitive information
    masked_repair_password = '****' if repair_password else None
    masked_repair_username = '****' if repair_username else None
    masked_repair_encrypt_recovery_key = '****' if encrypt_recovery_key else None
    logger.debug('vm repair create command parameters: vm_name: %s, resource_group_name: %s, repair_password: %s, repair_username: %s, repair_vm_name: %s, copy_disk_name: %s, repair_group_name: %s, unlock_encrypted_vm: %s, enable_nested: %s, associate_public_ip: %s, distro: %s, encrypt_recovery_key: %s, disable_trusted_launch: %s, os_disk_type: %s, tags: %s, copy_tags: %s, size: %s',
                 vm_name, resource_group_name, masked_repair_password, masked_repair_username, repair_vm_name, copy_disk_name, repair_group_name, unlock_encrypted_vm, enable_nested, associate_public_ip, distro, masked_repair_encrypt_recovery_key, disable_trusted_launch, os_disk_type, tags, copy_tags, size)

    # Initializing a command helper object.
    command = command_helper(logger, cmd, 'vm repair create')

    # The main command execution block.
    try:
        # Deprecation warning for --yes parameter
        if yes:
            logger.warning('The --yes parameter is deprecated and will be removed in a future release. There is currently no behavior change associated with this parameter.')

        # TODO: add permissions checks - can user create VMs, disks, resource groups, etc.
        # Set parameters used in exception handling to avoid Unbound errors:
        existing_rg = None
        copy_disk_id = None

        # Fetching the data of the source VM.
        source_vm = get_vm(cmd, resource_group_name, vm_name)
        source_vm_instance_view = get_vm(cmd, resource_group_name, vm_name, 'instanceView')

        # Checking if the OS of the source VM is Linux and what the Hyper-V generation is.
        is_linux = _is_linux_os(source_vm)
        vm_hypervgen = _is_gen2(source_vm_instance_view)
        source_controller = _fetch_source_disk_controller_type(source_vm)
        _set_source_resource_context(command, source_vm, source_vm_instance_view,
                                     disk_controller_type=source_controller)

        # Fetching the name of the OS disk and checking if it's managed.
        target_disk_name = source_vm.storage_profile.os_disk.name
        is_managed = _uses_managed_disk(source_vm)

        # Set up tags variable with passed data and resource.  Passed variable 'merged_tags' will be the holding location for the data throughout.
        merged_tags = {}
        # Optionally copy existing VM tags from the source VM.
        if copy_tags and source_vm.tags:
            merged_tags.update(source_vm.tags)
        # Merge user-provided tags
        if isinstance(tags, dict):
            merged_tags.update(tags)
        elif tags:
            # Azure CLI supports space-separated key=value strings
            for item in str(tags).split():
                # split the item into key and value at the first '=', we'll ignore items that don't have both key and value as invalid
                key, sep, value = item.partition('=')
                if sep:
                    merged_tags[key] = value

        # Add the generated repair resource tag
        repair_key, sep, repair_value = _get_repair_resource_tag(
            resource_group_name, vm_name
        ).partition('=')
        if sep:
            merged_tags[repair_key] = repair_value

        # Validate tag keys and values before they are placed into the command string.
        # Tag values can originate from the source VM (via --copy-tags) and are therefore
        # untrusted. _validate_tags_for_command rejects double quotes, control characters and
        # the cmd.exe expansion characters '%' and '!' (which cannot be safely escaped on a
        # 'cmd /c' command line); every other character, including shell metacharacters such
        # as & | < >, is preserved and passed through literally. See MSRC 115198 / VULN-185362.
        _validate_tags_for_command(merged_tags)

        # Convert to CLI string for passing to az cli later.
        # Each key=value token is quoted with shlex.quote so values containing spaces or
        # shell metacharacters survive re-tokenization in _call_az_command as a single
        # argument. Combined with the Windows cmd.exe quoting in _call_az_command, this
        # prevents command injection through attacker-controlled source VM tags.
        # See MSRC 115198 / VULN-185362.
        tag_string = ' '.join(shlex.quote(f'{tag_key}={tag_value}') for tag_key, tag_value in merged_tags.items())

        # initializing the list of created resources.
        created_resources = []
        # Fetching the architecture of the source VM.
        architecture_type = _fetch_architecture(source_vm)

        # Checking if the source VM's OS is Linux and if it uses a managed disk.
        if is_linux and _uses_managed_disk(source_vm):
            # Setting the OS type to 'Linux'.
            os_type = 'Linux'
            # Checking the Hyper-V generation of the source VM.
            hyperV_generation_linux = _check_linux_hyperV_gen(source_vm)
            if hyperV_generation_linux == 'V2':
                # If the Hyper-V generation is 'V2', it may be ARM:
                if architecture_type == 'Arm64':
                    # If the architecture type is 'Arm64', log this information and select the Linux distribution for an Arm64 VM.
                    logger.info('ARM64 VM detected')
                    os_image_urn = _select_distro_linux_Arm64(distro)
                    # Trusted launch is not supported on ARM
                    logger.info('Disabling trusted launch on ARM')
                    disable_trusted_launch = True
                else:
                    # log this information and select the Linux distribution for an x86 Gen2 VM.
                    logger.info('Generation 2 VM detected')
                    os_image_urn = _select_distro_linux_gen2(distro)
            else:
                # If the architecture type is not 'V2', select a Gen1 VM
                os_image_urn = _select_distro_linux(distro)
        else:
            # TODO: This logic should be improved - the checks in the called functions are very simple and defaults are outdated - Windows SME input required
            # If the source VM's OS is not Linux, check if a recovery key is provided.
            if encrypt_recovery_key:
                # If a recovery key is provided, fetch the same Windows OS URN for a VM with Bitlocker encryption.
                # this also covers Trusted Launch VMs which may have additional security needs, so use the same image
                os_image_urn = _fetch_matching_windows_os_urn(source_vm)
            else:
                # If no recovery key is provided, fetch the compatible Windows OS URN for a regular VM.
                os_image_urn = _fetch_compatible_windows_os_urn(source_vm, source_vm_instance_view)
            # Setting the OS type to 'Windows'.
            os_type = 'Windows'

        # Set public IP address for repair VM
        public_ip_name = _make_public_ip_name(repair_vm_name, associate_public_ip)

        # Set up base create vm command
        # Always include the --public-ip-address argument to 'vm create', but only provide a value if associate_public_ip is True.
        # Old code would omit the argument entirely if False, which leads to creating a public IP implicitly in some scenarios.  The desired behavior
        # is to never create a public IP unless explicitly requested.
        if associate_public_ip:
            public_ip_arg = f' --public-ip-address {public_ip_name}'
        else:
            public_ip_arg = ' --public-ip-address ""'

        # Assemble the create VM command for passing to azure cli call
        create_repair_vm_command = (
            f'az vm create -g {repair_group_name} -n {repair_vm_name} '
            f'--image {os_image_urn} '
            f'--admin-username {repair_username} '
            f'--admin-password {repair_password}'
            f'{public_ip_arg} '
            f'--tags {tag_string}'
        )
        # Linux-specific consideration(s)
        if is_linux:
            create_repair_vm_command += f' --custom-data {_get_cloud_init_script()}'

        # Fetching the size of the repair VM.
        sku = _fetch_compatible_sku(source_vm, enable_nested, size)

        if not sku:
            # If no compatible size is found, raise an error.
            raise SkuNotAvailableError('Failed to find compatible VM size for source VM\'s OS disk within given region and subscription.')
        # Adding the size to the command.
        create_repair_vm_command += ' --size {sku}'.format(sku=sku)

        supported_controllers = []
        if not disk_controller_type and source_controller and str(source_controller).lower() == 'nvme':
            supported_controllers = _fetch_sku_disk_controller_types(sku, source_vm.location)
        selected_controller, level, message = _select_repair_disk_controller_type(
            source_controller, supported_controllers, disk_controller_type)
        getattr(logger, level)(message)
        if level == 'warning':
            command.warnings.append(message)
        if selected_controller:
            create_repair_vm_command += ' --disk-controller-type {controller}'.format(controller=selected_controller)
        if hasattr(command, 'set_resource_context'):
            command.set_resource_context(repair_vm_disk_controller_type=selected_controller)

        # Setting the availability zone for the repair VM.
        # If the source VM has availability zones, the first one is chosen for the repair VM.
        if source_vm.zones:
            zone = source_vm.zones[0]
            create_repair_vm_command += ' --zone {zone}'.format(zone=zone)

        if disable_trusted_launch:
            logger.debug('Set security-type to Standard...')
            create_repair_vm_command += ' --security-type Standard'
        else:
            # If a Bitlocker recovery key is provided, this indicates the source VM is encrypted.
            # In this case, the VM and OS disk security profiles need to be fetched and added to the repair VM creation command.
            if encrypt_recovery_key:
                # TODO: this was assumed to also need for Trusted Launch VMs, but I don't think this is the case.

                # For confidential VM, some SKUs expect specific security types, secure_boot_enabled and vtpm_enabled.
                # Fetching the VM security profile and adding it to the command if it exists.
                logger.debug('Fetching VM security profile...')
                vm_security_params = _fetch_vm_security_profile_parameters(source_vm)
                if vm_security_params:
                    create_repair_vm_command += vm_security_params

        if encrypt_recovery_key:
            # TODO: this was assumed to also need for Trusted Launch VMs, but I don't think this is the case.

            # For confidential VM and Trusted Launch VM security tags on disks, the disk security profile needs to be brought over as well.
            # Fetching the OS Disk security profile and adding it to the command if it exists.
            logger.debug('Fetching OS Disk security profile...')
            osdisk_security_params = _fetch_osdisk_security_profile_parameters(source_vm)
            if osdisk_security_params:
                create_repair_vm_command += osdisk_security_params

        # Creating a new resource group for the repair VM and its resources.
        # First, check if the repair group already exists.
        # If it doesn't, create a new resource group at the same location as the source VM.
        existing_rg = _check_existing_rg(repair_group_name)
        if not existing_rg:
            create_resource_group_command = 'az group create -l {loc} -n {group_name}' \
                .format(loc=source_vm.location, group_name=repair_group_name)
            logger.info('Creating resource group for repair VM and its resources...')
            _call_az_command(create_resource_group_command)

        # Check if user is changing the Repair VM os disk type
        if os_disk_type:
            create_repair_vm_command += ' --storage-sku {os_disk_type} '.format(os_disk_type=os_disk_type)

        # Check if the source VM uses managed disks.
        # If it does, the repair VM will also be created with managed disks.
        if is_managed:
            logger.info('Source VM uses managed disks. Creating repair VM with managed disks.\n')

            # Fetch the SKU, location, OS type, and Hyper-V generation of the disk from the source VM.
            disk_sku, location, os_type, hyperV_generation, _ = _fetch_disk_info(resource_group_name, target_disk_name)

            # Prepare the command to create a copy of the source VM's OS disk.
            # The command includes the resource group name, copy disk name, target disk name, SKU, location, and OS type.
            copy_disk_command = 'az disk create -g {g} -n {n} --source {s} --sku {sku} --location {loc} --os-type {os_type} --query id -o tsv' \
                .format(g=resource_group_name, n=copy_disk_name, s=target_disk_name, sku=disk_sku, loc=location, os_type=os_type)
            # If the Hyper-V generation for the disk is available, append it to the copy disk command.
            if hyperV_generation:
                copy_disk_command += ' --hyper-v-generation {hyperV}'.format(hyperV=hyperV_generation)

            # If the source VM is a Linux Gen2 VM but the Hyper-V generation is not available in the disk info,
            # log this situation and manually add 'V2' to the copy disk command.
            elif is_linux and hyperV_generation_linux == 'V2':
                logger.info('The disk did not contain the information of gen2, but the machine is created from gen2 image')
                copy_disk_command += ' --hyper-v-generation {hyperV}'.format(hyperV=hyperV_generation_linux)

            # If the source VM has availability zones, get the first one and add it to the copy disk command.
            if source_vm.zones:
                zone = source_vm.zones[0]
                copy_disk_command += ' --zone {zone}'.format(zone=zone)

            # Execute the command to create a copy of the OS disk of the source VM.
            logger.info('Copying OS disk of source VM...')
            copy_disk_id = _call_az_command(copy_disk_command).strip('\n')

            # Depending on the operating system of the source VM and whether it's encrypted, different steps are taken.
            # If the source VM is not a Linux machine, create the repair VM.
            # This is the case for Windows VMs, both encrypted and not encrypted.
            if not is_linux:
                # Call the method to create the repair VM, providing the necessary parameters.
                _create_repair_vm(copy_disk_id, create_repair_vm_command, repair_password, repair_username)

            # If the source VM is a Windows machine and it is encrypted, unlock the encrypted VM after creation.
            if not is_linux and unlock_encrypted_vm:
                # Call the method to create the repair VM.
                _create_repair_vm(copy_disk_id, create_repair_vm_command, repair_password, repair_username)
                # Call the method to unlock the encrypted VM, providing the necessary parameters.
                _unlock_encrypted_vm_run(repair_vm_name, repair_group_name, is_linux, encrypt_recovery_key)

            # If the source VM is a Linux machine and it is encrypted, create the repair VM and then unlock it.
            if is_linux and unlock_encrypted_vm:
                # Call the method to create the repair VM.
                _create_repair_vm(copy_disk_id, create_repair_vm_command, repair_password, repair_username)
                # Call the method to unlock the encrypted VM.
                _unlock_encrypted_vm_run(repair_vm_name, repair_group_name, is_linux)

            # If the source VM is a Linux machine and it is not encrypted, create the repair VM and then attach the data disk.
            # This is done after VM creation to avoid a UUID mismatch causing an incorrect boot.
            if is_linux and (not unlock_encrypted_vm):
                # Call the method to create the repair VM, with the fix_uuid parameter set to True.
                _create_repair_vm(copy_disk_id, create_repair_vm_command, repair_password, repair_username, fix_uuid=True)
                logger.info('Attaching copied disk to repair VM as data disk...')
                # Set up the command to attach the copied disk to the repair VM.
                attach_disk_command = "az vm disk attach -g {g} --name {disk_id} --vm-name {vm_name} ".format(g=repair_group_name, disk_id=copy_disk_id, vm_name=repair_vm_name)
                # Execute the command to attach the disk.
                _call_az_command(attach_disk_command)

        # Check if the source VM uses unmanaged disks.
        # If it does, the repair VM will also be created with unmanaged disks.
        else:
            logger.info('Source VM uses unmanaged disks. Creating repair VM with unmanaged disks.\n')

            # Get the URI of the OS disk from the source VM.
            os_disk_uri = source_vm.storage_profile.os_disk.vhd.uri

            # Create the name of the copy disk by appending '.vhd' to the existing name.
            copy_disk_name = copy_disk_name + '.vhd'

            # Create a StorageResourceIdentifier object for the storage account associated with the OS disk.
            storage_account = StorageResourceIdentifier(cmd.cli_ctx.cloud, os_disk_uri)

            # Validate the VM creation command to ensure all parameters are correct before proceeding.
            validate_create_vm_command = create_repair_vm_command + ' --validate'
            logger.info('Validating VM template before continuing...')
            _call_az_command(validate_create_vm_command, secure_params=[repair_password, repair_username])

            # Fetch the connection string of the storage account.
            get_connection_string_command = 'az storage account show-connection-string -g {g} -n {n} --query connectionString -o tsv' \
                .format(g=resource_group_name, n=storage_account.account_name)
            logger.debug('Fetching storage account connection string...')
            connection_string = _call_az_command(get_connection_string_command).strip('\n')

            # Create a snapshot of the unmanaged OS disk.
            make_snapshot_command = 'az storage blob snapshot -c {c} -n {n} --connection-string "{con_string}" --query snapshot -o tsv' \
                .format(c=storage_account.container, n=storage_account.blob, con_string=connection_string)
            logger.info('Creating snapshot of OS disk...')
            snapshot_timestamp = _call_az_command(make_snapshot_command, secure_params=[connection_string]).strip('\n')
            snapshot_uri = os_disk_uri + '?snapshot={timestamp}'.format(timestamp=snapshot_timestamp)

            # Create a copy of the snapshot into an unmanaged disk.
            copy_snapshot_command = 'az storage blob copy start -c {c} -b {name} --source-uri {source} --connection-string "{con_string}"' \
                .format(c=storage_account.container, name=copy_disk_name, source=snapshot_uri, con_string=connection_string)
            logger.info('Creating a copy disk from the snapshot...')
            _call_az_command(copy_snapshot_command, secure_params=[connection_string])

            # Generate the URI of the copied disk.
            copy_disk_id = os_disk_uri.rstrip(storage_account.blob) + copy_disk_name

            # Create the repair VM with the copied unmanaged disk.
            create_repair_vm_command = create_repair_vm_command + ' --use-unmanaged-disk'
            logger.info('Creating repair VM while disk copy is in progress...')
            _call_az_command(create_repair_vm_command, secure_params=[repair_password, repair_username])

            # Check if the disk copy process is done.
            logger.info('Checking if disk copy is done...')
            copy_check_command = 'az storage blob show -c {c} -n {name} --connection-string "{con_string}" --query properties.copy.status -o tsv' \
                .format(c=storage_account.container, name=copy_disk_name, con_string=connection_string)
            copy_result = _call_az_command(copy_check_command, secure_params=[connection_string]).strip('\n')

            # If the disk copy process failed, raise an error.
            if copy_result != 'success':
                raise UnmanagedDiskCopyError('Unmanaged disk copy failed.')

            # Attach the copied unmanaged disk to the repair VM.
            logger.info('Attaching copied disk to repair VM as data disk...')
            attach_disk_command = "az vm unmanaged-disk attach -g {g} -n {disk_name} --vm-name {vm_name} --vhd-uri {uri}" \
                .format(g=repair_group_name, disk_name=copy_disk_name, vm_name=repair_vm_name, uri=copy_disk_id)
            _call_az_command(attach_disk_command)

        # Check if the Nested Hyper-V needs to be enabled.
        # If it does, run the script to install Hyper-V and create the nested VM.
        if enable_nested:
            logger.info("Running Script win-enable-nested-hyperv.ps1 to install HyperV")

            # Set up the command to run the script to enable Nested Hyper-V.
            run_hyperv_command = "az vm repair run -g {g} -n {name} --run-id win-enable-nested-hyperv --parameters gen={gen}" \
                .format(g=repair_group_name, name=repair_vm_name, gen=vm_hypervgen)

            # Execute the command to enable Nested Hyper-V.
            ret_enable_nested = _call_az_command(run_hyperv_command)
            logger.debug("az vm repair run hyperv command returned: %s", ret_enable_nested)

            # If the script indicates that a restart is required, restart the repair VM.
            if str.find(ret_enable_nested, "SuccessRestartRequired") > -1:
                restart_cmd = 'az vm restart -g {rg} -n {vm}'.format(rg=repair_group_name, vm=repair_vm_name)
                logger.info("Restarting Repair VM")
                restart_ret = _call_az_command(restart_cmd)
                logger.debug(restart_ret)

                # After the restart, run the script to enable Nested Hyper-V again to create the nested VM.
                logger.info("Running win-enable-nested-hyperv.ps1 again to create nested VM")
                run_hyperv_command = "az vm repair run -g {g} -n {name} --run-id win-enable-nested-hyperv --parameters gen={gen}" \
                    .format(g=repair_group_name, name=repair_vm_name, gen=vm_hypervgen)
                ret_enable_nested_again = _call_az_command(run_hyperv_command)
                logger.debug("stderr: %s", ret_enable_nested_again)

        # List all the resources in the repair resource group.
        created_resources = _list_resource_ids_in_rg(repair_group_name)

        # Set the command status to success.
        command.set_status_success()

    # Some error happened. Stop command and clean-up resources.
    except KeyboardInterrupt:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = "Command interrupted by user input."
        command.message = "Command interrupted by user input. Cleaning up resources."
    except AzCommandError as azCommandError:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(azCommandError)
        command.message = "Repair create failed. Cleaning up created resources."
    except SkuDoesNotSupportHyperV as skuDoesNotSupportHyperV:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(skuDoesNotSupportHyperV)
        command.message = "provided sku does not support nested VM in hyperv. Please run command without --enabled-nested or provide a valid --size parameter. Cleaning up created resources."
    except ScriptReturnsError as scriptReturnsError:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(scriptReturnsError)
        command.message = "Error returned from script when enabling hyperv."
    except SkuNotAvailableError as skuNotAvailableError:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(skuNotAvailableError)
        command.message = "Please check if the current subscription can create more VM resources. Cleaning up created resources."
    except UnmanagedDiskCopyError as unmanagedDiskCopyError:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(unmanagedDiskCopyError)
        command.message = "Repair create failed. Please try again at another time. Cleaning up created resources."
    except WindowsOsNotAvailableError:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = 'Compatible Windows OS image not available.'
        command.message = 'A compatible Windows OS image is not available at this time, please check subscription.'
    except Exception as exception:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(exception)
        command.message = 'An unexpected error occurred. Try running again with the --debug flag to debug.'

    finally:
        if command.error_stack_trace:
            logger.debug(command.error_stack_trace)
    # Generate return results depending on command state
    if not command.is_status_success():
        command.set_status_error()
        return_dict = command.init_return_dict()
        if existing_rg:
            _clean_up_resources(repair_group_name, confirm=True)
        else:
            _clean_up_resources(repair_group_name, confirm=False)
    else:
        created_resources.append(copy_disk_id if copy_disk_id is not None else "")
        command.message = 'Your repair VM \'{n}\' has been created in the resource group \'{repair_rg}\' with disk \'{d}\' attached as data disk. ' \
                          'Please use this VM to troubleshoot and repair. Once the repairs are complete use the command ' \
                          '\'az vm repair restore -n {source_vm} -g {rg} --verbose\' to restore disk to the source VM. ' \
                          'Note that the copied disk is created within the original resource group \'{rg}\'.' \
                          .format(n=repair_vm_name, repair_rg=repair_group_name, d=copy_disk_name, rg=resource_group_name, source_vm=vm_name)
        return_dict = command.init_return_dict()
        # Add additional custom return properties
        return_dict['repair_vm_name'] = repair_vm_name
        return_dict['copied_disk_name'] = copy_disk_name
        return_dict['copied_disk_uri'] = copy_disk_id if copy_disk_id is not None else ""
        return_dict['repair_resource_group'] = repair_group_name
        return_dict['created_resources'] = created_resources

        logger.info('\n%s\n', command.message)
    return return_dict


# This method is responsible for restoring the VM after repair
def restore(cmd, vm_name, resource_group_name, disk_name=None, repair_vm_id=None, yes=False, no_cleanup=False):

    # Create an instance of the command helper object to facilitate logging and status tracking.
    command = command_helper(logger, cmd, 'vm repair restore')
    source_disk = None
    repair_resource_group = None

    try:
        # Fetch source and repair VM data
        source_vm = get_vm(cmd, resource_group_name, vm_name)  # Fetch the source VM data
        _enrich_source_resource_context(
            command, cmd, source_vm, resource_group_name, vm_name)
        is_managed = _uses_managed_disk(source_vm)  # Check if the source VM uses managed disks
        if repair_vm_id:
            logger.info('Repair VM ID: %s', repair_vm_id)
            repair_vm_id = parse_resource_id(repair_vm_id)  # Parse the repair VM ID
            repair_vm_name = repair_vm_id['name']
            repair_resource_group = repair_vm_id['resource_group']
            _enrich_repair_controller_context(
                command, cmd, repair_resource_group, repair_vm_name)

            # For MANAGED DISK
            if is_managed:
                source_disk = source_vm.storage_profile.os_disk.name
                # Retrieve required data from the repair disk before doing the disk swap, we need the full disk URI (id), not just the name
                _, _, _, _, disk_id = _fetch_disk_info(resource_group_name, disk_name)
                # Commands to detach the repaired data disk from the repair VM and os-disk-swap it onto the source VM
                detach_disk_command = 'az vm disk detach -g {g} --vm-name {repair} --name {disk}' \
                    .format(g=repair_resource_group, repair=repair_vm_name, disk=disk_name)
                attach_fixed_command = 'az vm update -g {g} -n {n} --os-disk {disk}' \
                    .format(g=resource_group_name, n=vm_name, disk=disk_id)

                # Detach the repaired data disk from the repair VM and attach it to the source VM as an OS disk
                logger.info('Detaching repaired data disk from repair VM...')
                _call_az_command(detach_disk_command)
                logger.info('Attaching repaired data disk to source VM as an OS disk...')
                _call_az_command(attach_fixed_command)

            # For UNMANAGED DISK
            else:
                source_disk = source_vm.storage_profile.os_disk.vhd.uri
                # Fetch disk uri from disk name
                repair_vm = get_vm(cmd, repair_vm_id['resource_group'], repair_vm_id['name'])
                data_disks = repair_vm.storage_profile.data_disks

                # The params went through validator so no need for existence checks
                disk_uri = [disk.vhd.uri for disk in data_disks if disk.name == disk_name][0]

                # Commands to detach the repaired data disk from the repair VM and attach it to the source VM as an OS disk
                detach_unamanged_command = 'az vm unmanaged-disk detach -g {g} --vm-name {repair} --name {disk}' \
                    .format(g=repair_resource_group, repair=repair_vm_name, disk=disk_name)
                attach_unmanaged_command = 'az vm update -g {g} -n {n} --set storageProfile.osDisk.vhd.uri="{uri}"' \
                    .format(g=resource_group_name, n=vm_name, uri=disk_uri)

                # Detach the repaired data disk from the repair VM and attach it to the source VM as an OS disk
                logger.info('Detaching repaired data disk from repair VM...')
                _call_az_command(detach_unamanged_command)
                logger.info('Attaching repaired data disk to source VM as an OS disk...')
                _call_az_command(attach_unmanaged_command)

            # Clean up the resources in the repair resource group
            _clean_up_resources(repair_resource_group, confirm=not yes, skip_cleanup=no_cleanup)
            command.set_status_success()  # Set the command status to success
    # Handle possible exceptions
    except KeyboardInterrupt:
        # Capture the stack trace and set the error message if the command is interrupted by the user
        command.error_stack_trace = traceback.format_exc()
        command.error_message = "Command interrupted by user input."
        command.message = "Command interrupted by user input. If the restore command fails at retry, please rerun the repair process from \'az vm repair create\'."
    except AzCommandError as azCommandError:
        # Capture the stack trace and set the error message if an Azure command error occurs
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(azCommandError)
        command.message = "Repair restore failed. If the restore command fails at retry, please rerun the repair process from \'az vm repair create\'."
    except Exception as exception:
        # Capture the stack trace and set the error message if an unexpected error occurs
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(exception)
        command.message = 'An unexpected error occurred. Try running again with the --debug flag to debug.'
    finally:
        # Log the stack trace if an error has occurred
        if command.error_stack_trace:
            logger.debug(command.error_stack_trace)

    # Set the command status to error if it's not success
    if not command.is_status_success():
        command.set_status_error()
        return_dict = command.init_return_dict()
    else:
        # Construct return dict
        command.message = '\'{disk}\' successfully attached to \'{n}\' as an OS disk. Please test your repairs and once confirmed, ' \
            'you may choose to delete the source OS disk \'{src_disk}\' within resource group \'{rg}\' manually if you no longer need it, to avoid any undesired costs.' \
            .format(disk=disk_name, n=vm_name, src_disk=source_disk, rg=resource_group_name)
        if no_cleanup and repair_resource_group:
            command.message += ' The repair resources in the resource group \'{repair_rg}\' were kept because --no-cleanup was used. ' \
                'Delete them with \'az group delete --name {repair_rg}\' once you no longer need them, to avoid any undesired costs.' \
                .format(repair_rg=repair_resource_group)
        return_dict = command.init_return_dict()
        logger.info('\n%s\n', return_dict['message'])

    return return_dict


# This method is responsible for running a script on a VM
def run(cmd, vm_name, resource_group_name, run_id=None, repair_vm_id=None, custom_script_file=None, parameters=None, run_on_repair=False, preview=None):

    # Log the input parameters for the method
    logger.debug('vm repair run parameters: vm_name: %s, resource_group_name: %s, run_id: %s, repair_vm_id: %s, custom_script_file: %s, parameters: %s, run_on_repair: %s, preview: %s',
                 vm_name, resource_group_name, run_id, repair_vm_id, custom_script_file, parameters, run_on_repair, preview)

    # Reject a bad preview url before the command helper exists: its destructor runs at interpreter
    # shutdown when the command aborts early, which loses the telemetry and prints a shutdown traceback.
    if preview:
        _parse_preview_url(preview)

    # Initiate a command helper object for logging and status tracking
    command = command_helper(logger, cmd, 'vm repair run')

    # Define the script names for Linux and Windows
    LINUX_RUN_SCRIPT_NAME = 'linux-run-driver.sh'
    WINDOWS_RUN_SCRIPT_NAME = 'win-run-driver.ps1'

    if preview:
        _set_repair_map_url(preview)

    try:
        # Fetch data of the VM on which the script is to be run
        source_vm = get_vm(cmd, resource_group_name, vm_name)

        # Determine the OS of the source VM
        is_linux = _is_linux_os(source_vm)
        _enrich_source_resource_context(
            command, cmd, source_vm, resource_group_name, vm_name)

        # Choose the appropriate script based on the OS of the source VM
        if is_linux:
            script_name = LINUX_RUN_SCRIPT_NAME
        else:
            script_name = WINDOWS_RUN_SCRIPT_NAME

        # If run_on_repair is False, then the repair VM is the same as the source VM (i.e., scripts run directly on the source VM)
        if run_on_repair:
            repair_vm_id = parse_resource_id(repair_vm_id)
            repair_vm_name = repair_vm_id['name']
            repair_resource_group = repair_vm_id['resource_group']
            _enrich_repair_controller_context(
                command, cmd, repair_resource_group, repair_vm_name)
        else:
            repair_vm_name = vm_name
            repair_resource_group = resource_group_name
            if hasattr(command, 'set_resource_context'):
                command.set_resource_context(
                    repair_vm_disk_controller_type=command.disk_controller_type)

        run_command_params = []
        additional_scripts = []

        # For the default scenario where a run ID is provided
        if not custom_script_file:
            # Fetch the path to the script from GitHub using the run ID
            repair_script_path = _fetch_run_script_path(run_id)
            run_command_params.append('script_path="./{}"'.format(repair_script_path))
        # For the custom script scenario for script testers
        else:
            run_command_params.append('script_path=no-op')
            additional_scripts.append(custom_script_file)

        # If a preview URL is provided, validate it and extract the fork and branch names
        run_command_params.extend(_build_repo_params(preview, is_linux))

        # Append parameters for the script
        if parameters:
            if is_linux:
                param_string = _process_bash_parameters(parameters)
            else:
                param_string = _process_ps_parameters(parameters)
            run_command_params.append('params="{}"'.format(param_string))

        if run_on_repair:
            vm_string = 'repair VM'
        else:
            vm_string = 'VM'

        logger.info('Running script on %s: %s', vm_string, repair_vm_name)  # Log the VM on which the script is being run

        # Start the timer to measure the run-time of the script
        script_start_time = timeit.default_timer()

        # Invoke the run command on the VM and capture the standard output and error
        stdout, stderr = _invoke_run_command(script_name, repair_vm_name, repair_resource_group, is_linux, run_command_params, additional_scripts)

        # Calculate the run-time of the script and log it
        command.script.run_time = timeit.default_timer() - script_start_time
        logger.debug("stderr: %s", stderr)

        # Parse the standard output to check if the script execution was successful
        run_script_succeeded = _check_script_succeeded(stdout)

        # Parse the raw logs from the standard output
        logs = _parse_run_script_raw_logs(stdout)

        # Process the start and end of the log
        # If the log is over 4k bytes, it gets cutoff at the start
        log_cutoff = True
        log_fullpath = ''
        for log in logs:
            if log['level'] == 'Log-Start':
                log_cutoff = False
            if log['level'] == 'Log-End':
                split_log = log['message'].split(']')
                if len(split_log) == 2:
                    log_fullpath = split_log[1]

        # If the log is cutoff, give a warning to the user
        if log_cutoff:
            logger.warning('Log file is too large and has been cutoff at the start of file. Please locate the log file within the %s using the logFullpath to check full logs.', vm_string)

        # If the script execution was successful, set the status to success and log the output
        # If the script execution was unsuccessful, set the status to error and log the error
        if run_script_succeeded:
            command.script.set_status_success()
            command.message = 'Script completed successfully.'
            command.script.output = '\n'.join([log['message'] for log in logs if log['level'].lower() == 'output'])
            logger.info('\nScript returned with output:\n%s\n', command.script.output)
        else:
            command.script.set_status_error()
            command.message = 'Script completed with errors.'
            command.script.output = '\n'.join([log['message'] for log in logs if log['level'].lower() == 'error'])
            logger.error('\nScript returned with error:\n%s\n', command.script.output)

        # Set the overall command status to success
        command.set_status_success()

    except KeyboardInterrupt:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = "Command interrupted by user input."
        command.message = "Repair run failed. Command interrupted by user input."
    except AzCommandError as azCommandError:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(azCommandError)
        command.message = "Repair run failed."
    except requests.exceptions.RequestException as exception:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(exception)
        command.message = "Failed to fetch run script data from GitHub. Please check this repository is reachable: https://github.com/Azure/repair-script-library"
    except RunScriptNotFoundForIdError as exception:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(exception)
        command.message = "Repair run failed. Run ID not found."
    except Exception as exception:
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(exception)
        command.message = 'An unexpected error occurred. Try running again with the --debug flag to debug.'
    finally:
        if command.error_stack_trace:
            logger.debug(command.error_stack_trace)

    if not command.is_status_success():
        command.set_status_error()
        command.script.output = 'Repair run failed.'
        return_dict = command.init_return_dict()
    else:
        # Build return Dict
        return_dict = command.init_return_dict()
        return_dict['script_status'] = command.script.status
        return_dict['logs'] = stdout
        return_dict['err'] = stderr
        return_dict['log_full_path'] = log_fullpath
        return_dict['output'] = command.script.output
        return_dict['vm_name'] = repair_vm_name
        return_dict['resource_group'] = repair_resource_group

    return return_dict


# This method lists all available repair scripts
def list_scripts(cmd, preview=None):
    # Initiate a command helper object for logging and status tracking
    if preview:
        _parse_preview_url(preview)

    command = command_helper(logger, cmd, 'vm repair list-scripts')

    if preview:
        _set_repair_map_url(preview)

    try:
        # Fetch the map of all available scripts
        run_map = _fetch_run_script_map()

        # Set the command status to success
        command.set_status_success()

    # Handle possible exceptions
    except requests.exceptions.RequestException as exception:
        # Capture the stack trace and set the error message if a request exception occurs
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(exception)
        command.message = "Failed to fetch run script data from GitHub. Please check this repository is reachable: https://github.com/Azure/repair-script-library"
    except Exception as exception:
        # Capture the stack trace and set the error message if an unexpected error occurs
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(exception)
        command.message = 'An unexpected error occurred. Try running again with the --debug flag to debug.'
    finally:
        # Log the stack trace if an error has occurred
        if command.error_stack_trace:
            logger.debug(command.error_stack_trace)

    # Set the command status to error if it's not success
    if not command.is_status_success():
        command.set_status_error()
        return_dict = command.init_return_dict()
    else:
        # Construct return dict
        command.message = 'Available script list successfully fetched from https://github.com/Azure/repair-script-library'
        return_dict = command.init_return_dict()
        return_dict['map'] = run_map

    return return_dict


# This method resets the network interface card (NIC) of a virtual machine (VM)
def reset_nic(cmd, vm_name, resource_group_name, yes=False):
    # Initiate a command helper object for logging and status tracking
    command = command_helper(logger, cmd, 'vm repair reset-nic')
    DYNAMIC_CONFIG = 'Dynamic'

    try:
        # 0) Check if VM is off. If it is, ask to start the VM.
        # VM must be running to reset its NIC.
        VM_OFF_MESSAGE = 'VM is not running. The VM must be in running to reset its NIC.\n'

        vm_instance_view = get_vm(cmd, resource_group_name, vm_name, 'instanceView')
        VM_started = _check_n_start_vm(vm_name, resource_group_name, not yes, VM_OFF_MESSAGE, vm_instance_view)

        # If VM is not started, raise an error
        if not VM_started:
            raise CommandCanceledByUserError("Could not get consent to run VM before resetting the NIC.")

        # 1) Fetch VM network info
        logger.info('Fetching necessary VM network information to reset the NIC...\n')

        # Fetch primary NIC ID. The primary field is null or true for primary NICs.
        get_primary_nic_id_command = 'az vm nic list -g {g} --vm-name {n} --query "[[?primary].id || [?primary==null].id][0][0]" -o tsv' \
            .format(g=resource_group_name, n=vm_name)
        primary_nic_id = _call_az_command(get_primary_nic_id_command)

        # If no primary NIC is found, raise an error
        if not primary_nic_id:
            raise SupportingResourceNotFoundError('The primary NIC for the VM was not found on Azure.')
        primary_nic_name = primary_nic_id.split('/')[-1]

        # Get IP config info to get: vnet name, current private IP, IP config name, subnet ID
        get_primary_ip_config = 'az network nic ip-config list -g {g} --nic-name {nic_name} --query [[?primary]][0][0]' \
            .format(g=resource_group_name, nic_name=primary_nic_name)
        ip_config_string = _call_az_command(get_primary_ip_config)

        # If no primary IP configuration is found, raise an error
        if not ip_config_string:
            raise SupportingResourceNotFoundError('The primary IP configuration for the VM NIC was not found on Azure.')
        ip_config_object = json.loads(ip_config_string)

        # Extract subnet ID from the IP configuration
        subnet_id = ip_config_object['subnet']['id']

        # Tokenize the subnet ID to extract various components
        subnet_id_tokens = subnet_id.split('/')
        subnet_name = subnet_id_tokens[-1]  # Last token is subnet name
        vnet_name = subnet_id_tokens[-3]  # The third token from last is VNet name
        vnet_resource_group = subnet_id_tokens[-7]  # The seventh token from last is VNet resource group
        ipconfig_name = ip_config_object['name']  # Extract IP configuration name
        orig_ip_address = ip_config_object['privateIPAddress']  # Extract original IP address

        # Extract application security groups, if any
        application_names = ""
        applicationSecurityGroups = 'applicationSecurityGroups'
        if applicationSecurityGroups in ip_config_object:
            for item in ip_config_object[applicationSecurityGroups]:
                application_id_tokens = item['id'].split('/')
                if application_id_tokens[-1] is not None:
                    application_names += application_id_tokens[-1] + " "

        logger.info('applicationSecurityGroups {application_names}...\n')

        # Extract original IP allocation method (Dynamic or Static)
        orig_ip_allocation_method = ip_config_object['privateIPAllocationMethod']

        # Get available IP address within subnet
        get_available_ip_command = 'az network vnet subnet list-available-ips -g {g} --vnet-name {vnet} --name {subnet} --query [0] -o tsv' \
            .format(g=vnet_resource_group, vnet=vnet_name, subnet=subnet_name)
        swap_ip_address = _call_az_command(get_available_ip_command)
        if not swap_ip_address:
            # Raise available IP not found error
            raise SupportingResourceNotFoundError('Available IP address was not found within the VM subnet.')

        # 3) Update private IP address to another in subnet. This will invoke and wait for a VM restart.
        logger.info('Updating VM IP configuration. This might take a few minutes...\n')

        # Update IP address
        if application_names:
            update_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --private-ip-address {ip} --asgs {asgs}' \
                .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, ip=swap_ip_address, asgs=application_names)
        else:
            logger.info('applicationSecurityGroups do not exist...\n')
            update_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --private-ip-address {ip}' \
                .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, ip=swap_ip_address)
        _call_az_command(update_ip_command)

        # Wait for IP to be updated
        # This command waits for the network interface card (NIC) IP configuration to be updated
        wait_ip_update_command = 'az network nic ip-config wait --updated -g {g} --nic-name {nic}' \
            .format(g=resource_group_name, nic=primary_nic_name)
        _call_az_command(wait_ip_update_command)

        # 4) Revert the configurations. This will also invoke and wait for a VM restart.
        logger.info('NIC reset is complete. Now reverting back to your original configuration...\n')

        # Initialize the revert IP command variable
        revert_ip_command = None

        # If the original IP allocation method was dynamic, revert back to dynamic
        if orig_ip_allocation_method == DYNAMIC_CONFIG:
            # If there are application security groups, include them in the command
            if application_names:
                revert_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --set privateIpAllocationMethod={method} --asgs {asgs}' \
                    .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, method=DYNAMIC_CONFIG, asgs=application_names)
            else:
                revert_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --set privateIpAllocationMethod={method}' \
                    .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, method=DYNAMIC_CONFIG)
        else:
            # If the original IP allocation method was not dynamic, revert to the original static IP
            # If there are application security groups, include them in the command
            if application_names:
                revert_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --private-ip-address {ip} --asgs {asgs}' \
                    .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, ip=orig_ip_address, asgs=application_names)
            else:
                revert_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --private-ip-address {ip} ' \
                    .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, ip=orig_ip_address)

        # Execute the revert IP command
        _call_az_command(revert_ip_command)
        logger.info('VM guest NIC reset is complete and all configurations are reverted.')

    # Some error happened. Stop command and revert back as needed.
    except KeyboardInterrupt:
        command.set_status_error()
        command.error_stack_trace = traceback.format_exc()
        command.error_message = "Command interrupted by user input."
        command.message = "Command interrupted by user input."
    except AzCommandError as azCommandError:
        command.set_status_error()
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(azCommandError)
        command.message = "Reset NIC failed."
    except SupportingResourceNotFoundError as resourceError:
        command.set_status_error()
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(resourceError)
        command.message = "Reset NIC could not be initiated."
    except CommandCanceledByUserError as canceledError:
        command.set_status_error()
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(canceledError)
        command.message = VM_OFF_MESSAGE
    except Exception as exception:
        command.set_status_error()
        command.error_stack_trace = traceback.format_exc()
        command.error_message = str(exception)
        command.message = 'An unexpected error occurred. Try running again with the --debug flag to debug.'
    else:
        command.set_status_success()
        command.message = 'VM guest NIC reset complete. The VM is in running state.'
    finally:
        if command.error_stack_trace:
            logger.debug(command.error_stack_trace)
        # Generate return object and log errors if needed
        return_dict = command.init_return_dict()

    return return_dict


def repair_and_restore(cmd, vm_name, resource_group_name, repair_password=None, repair_username=None, repair_vm_name=None, copy_disk_name=None, repair_group_name=None, tags=None, copy_tags=False, size=None, no_cleanup=False):
    """
    This function manages the process of repairing and restoring a specified virtual machine (VM). The process involves
    the creation of a repair VM, the generation of a copy of the problem VM's disk, and the formation of a new resource
    group specifically for the repair operation.

    :param cmd: Command object that is used to execute Azure CLI commands.
    :param vm_name: The name of the VM that is targeted for repair.
    :param resource_group_name: The name of the resource group in which the targeted VM is located.
    :param repair_password: (Optional) Password to be used for the repair operation. If not provided, a random password is generated.
    :param repair_username: (Optional) Username to be used for the repair operation. If not provided, a random username is generated.
    :param repair_vm_name: (Optional) The name to assign to the repair VM. If not provided, a unique name is generated.
    :param copy_disk_name: (Optional) The name to assign to the copy of the disk. If not provided, a unique name is generated.
    :param repair_group_name: (Optional) The name of the repair resource group. If not provided, a unique name is generated.
    :param tags: (Optional) Tags to apply to the repair VM.
    :param copy_tags: (Optional) Boolean indicating whether to copy tags from the source VM to the repair VM.
    :param size: (Optional) The size of the repair VM.
    :param no_cleanup: (Optional) Boolean indicating whether the repair resources should be kept instead of deleted.
    """
    from datetime import datetime, timezone
    import secrets
    import string

    # Initialize command helper object
    command = command_helper(logger, cmd, 'vm repair repair-and-restore')

    try:
        source_vm = get_vm(cmd, resource_group_name, vm_name)
        _enrich_source_resource_context(
            command, cmd, source_vm, resource_group_name, vm_name)
    except Exception as exception:
        logger.debug('Could not fetch source VM resource shape for telemetry: %s', exception)

    # Generate a random password for the repair operation
    password_length = 30
    password_characters = string.ascii_lowercase + string.digits + string.ascii_uppercase
    repair_password = ''.join(secrets.choice(password_characters) for _ in range(password_length))

    # Generate a random username for the repair operation
    username_length = 20
    username_characters = string.ascii_lowercase + string.digits
    repair_username = ''.join(secrets.choice(username_characters) for _ in range(username_length))

    # Generate unique names for the repair VM, copied disk, and repair resource group
    timestamp = datetime.now(timezone.utc).strftime('%Y%m%d%H%M%S')
    repair_vm_name = ('repair-' + vm_name)[:14] + '_'
    copy_disk_name = vm_name + '-DiskCopy-' + timestamp
    repair_group_name = 'repair-' + vm_name + '-' + timestamp

    # Check if a resource group with the same name already exists
    existing_rg = _check_existing_rg(repair_group_name)

    # Create a repair VM, copy of the disk, and a new resource group
    create_out = create(cmd, vm_name, resource_group_name, repair_password, repair_username, repair_vm_name=repair_vm_name, copy_disk_name=copy_disk_name, repair_group_name=repair_group_name, associate_public_ip=False, tags=tags, copy_tags=copy_tags, size=size)

    # Log the output of the create operation
    logger.info('create_out: %s', create_out)

    # Extract the repair VM name, copied disk name, and repair resource group from the output
    repair_vm_name = create_out['repair_vm_name']
    copy_disk_name = create_out['copied_disk_name']
    repair_group_name = create_out['repair_resource_group']
    _enrich_repair_controller_context(
        command, cmd, repair_group_name, repair_vm_name)

    # Log that the fstab run command is about to be executed
    logger.info('Running fstab run command')

    try:
        # Run the fstab script on the repair VM
        run_out = run(cmd, repair_vm_name, repair_group_name, run_id='linux-alar2', parameters=["fstab", "initiator=SELFHELP"])

    except Exception:
        command.set_status_error()
        command.error_stack_trace = traceback.format_exc()
        command.error_message = "Command failed when running fstab script."
        command.message = "Command failed when running fstab script."

        # If the resource group existed before, confirm before cleaning up resources
        # Otherwise, clean up resources without confirmation
        if existing_rg:
            _clean_up_resources(repair_group_name, confirm=True, skip_cleanup=no_cleanup)
        else:
            _clean_up_resources(repair_group_name, confirm=False, skip_cleanup=no_cleanup)
        return

    # Log the output of the run command
    logger.info('run_out: %s', run_out)

    # If the fstab script returned an error, log the error and clean up resources
    if run_out['script_status'] == 'ERROR':
        logger.error('fstab script returned an error.')
        if existing_rg:
            _clean_up_resources(repair_group_name, confirm=True, skip_cleanup=no_cleanup)
        else:
            _clean_up_resources(repair_group_name, confirm=False, skip_cleanup=no_cleanup)
        return

    # Run the restore command
    logger.info('Running restore command')
    show_vm_id = 'az vm show -g {g} -n {n} --query id -o tsv' \
        .format(g=repair_group_name, n=repair_vm_name)

    repair_vm_id = _call_az_command(show_vm_id)

    restore(cmd, vm_name, resource_group_name, copy_disk_name, repair_vm_id, yes=True, no_cleanup=no_cleanup)

    # Set the success message
    repair_vm_fate = 'the repair resources were kept' if no_cleanup else 'the repair VM was then deleted'
    command.message = 'fstab script has been applied to the source VM. A new repair VM \'{n}\' was created in the resource group \'{repair_rg}\' with disk \'{d}\' attached as data disk. ' \
        'The repairs were complete using the fstab script and {fate}. ' \
        'The repair disk was restored to the source VM. ' \
        .format(n=repair_vm_name, repair_rg=repair_group_name, d=copy_disk_name, fate=repair_vm_fate)

    # Mark the operation as successful
    command.set_status_success()

    # If there were any errors, log the stack trace
    if command.error_stack_trace:
        logger.debug(command.error_stack_trace)

    # Initialize the return dictionary and log the success message
    return_dict = command.init_return_dict()

    logger.info('\n%s\n', command.message)

    # Return the result of the operation
    return return_dict


def repair_button(cmd, vm_name, resource_group_name, button_command, repair_password=None, repair_username=None, repair_vm_name=None, copy_disk_name=None, repair_group_name=None, tags=None, copy_tags=False, size=None, yes=False, no_cleanup=False):
    """
    Button-triggered repair operation. Supports tags for the repair VM.
    """
    from datetime import datetime, timezone
    import secrets
    import string

    # Init command helper object
    command = command_helper(logger, cmd, 'vm repair repair-button')

    # Deprecation warning for --yes parameter
    if yes:
        logger.warning('The --yes parameter is deprecated and will be removed in a future release. There is currently no behavior change associated with this parameter.')

    password_length = 30
    password_characters = string.ascii_lowercase + string.digits + string.ascii_uppercase
    repair_password = ''.join(secrets.choice(password_characters) for i in range(password_length))

    username_length = 20
    username_characters = string.ascii_lowercase + string.digits
    repair_username = ''.join(secrets.choice(username_characters) for i in range(username_length))

    timestamp = datetime.now(timezone.utc).strftime('%Y%m%d%H%M%S')
    repair_vm_name = ('repair-' + vm_name)[:14] + '_'
    copy_disk_name = vm_name + '-DiskCopy-' + timestamp
    repair_group_name = 'repair-' + vm_name + '-' + timestamp
    existing_rg = _check_existing_rg(repair_group_name)

    create_out = create(cmd, vm_name, resource_group_name, repair_password, repair_username, repair_vm_name=repair_vm_name, copy_disk_name=copy_disk_name, repair_group_name=repair_group_name, associate_public_ip=False, tags=tags, copy_tags=copy_tags, size=size)

    # log create_out
    logger.info('create_out: %s', create_out)

    repair_vm_name = create_out['repair_vm_name']
    copy_disk_name = create_out['copied_disk_name']
    repair_group_name = create_out['repair_resource_group']

    logger.info('Running command')

    try:
        run_out = run(cmd, repair_vm_name, repair_group_name, run_id='linux-alar2', parameters=[button_command, "initiator=SELFHELP"])

    except Exception:
        command.set_status_error()
        command.error_stack_trace = traceback.format_exc()
        command.error_message = "Command failed when running  script."
        command.message = "Command failed when running script."
        if existing_rg:
            _clean_up_resources(repair_group_name, confirm=True, skip_cleanup=no_cleanup)
        else:
            _clean_up_resources(repair_group_name, confirm=False, skip_cleanup=no_cleanup)
        return

    # log run_out
    logger.info('run_out: %s', run_out)

    if run_out['script_status'] == 'ERROR':
        logger.error(' script returned an error.')
        if existing_rg:
            _clean_up_resources(repair_group_name, confirm=True, skip_cleanup=no_cleanup)
        else:
            _clean_up_resources(repair_group_name, confirm=False, skip_cleanup=no_cleanup)
        return

    logger.info('Running restore command')
    show_vm_id = 'az vm show -g {g} -n {n} --query id -o tsv' \
        .format(g=repair_group_name, n=repair_vm_name)

    repair_vm_id = _call_az_command(show_vm_id)

    restore(cmd, vm_name, resource_group_name, copy_disk_name, repair_vm_id, yes=True, no_cleanup=no_cleanup)

    repair_vm_fate = 'the repair resources were kept' if no_cleanup else 'the repair VM was then deleted'
    command.message = 'script has been applied to the source VM. A new repair VM \'{n}\' was created in the resource group \'{repair_rg}\' with disk \'{d}\' attached as data disk. ' \
        'The repairs were complete using the script and {fate}. ' \
        'The repair disk was restored to the source VM. ' \
        .format(n=repair_vm_name, repair_rg=repair_group_name, d=copy_disk_name, fate=repair_vm_fate)

    command.set_status_success()
    if command.error_stack_trace:
        logger.debug(command.error_stack_trace)
    # Generate return object and log errors if needed
    return_dict = command.init_return_dict()

    logger.info('\n%s\n', command.message)

    return return_dict