[Vm-repair] Vm repair bug fixes

src/vm-repair/azext_vm_repair/_params.py

@@ -32,6 +32,7 @@ def load_arguments(self, _):
         c.argument('enable_nested', help='enable nested hyperv.')
         c.argument('associate_public_ip', help='Option to create repair vm with public ip')
         c.argument('distro', help='Option to create repair vm from a specific linux distro (rhel7|rhel8|suse12|ubuntu20|centos7|oracle7)')
+        c.argument('yes', help='Option to skip associate_public_ip prompt and confirm yes to it in no Tty mode')
 
     with self.argument_context('vm repair restore') as c:
         c.argument('repair_vm_id', help='Repair VM resource id.')

src/vm-repair/azext_vm_repair/_validators.py

@@ -87,7 +87,7 @@ def validate_create(cmd, namespace):
     # Validate vm password
     validate_vm_password(namespace.repair_password, is_linux)
     # Prompt input for public ip usage
-    if not namespace.associate_public_ip:
+    if (not namespace.associate_public_ip) and (not namespace.yes):
         _prompt_public_ip(namespace)
 
 
@@ -313,7 +313,7 @@ def fetch_repair_vm(namespace):
     # Find repair VM
     tag = _get_repair_resource_tag(namespace.resource_group_name, namespace.vm_name)
     try:
-        find_repair_command = 'az resource list --tag {tag} --query "[?type==\'microsoft.compute/virtualmachines\']" -o json' \
+        find_repair_command = 'az resource list --tag {tag} --query "[?type==\'microsoft.compute/virtualmachines\' || type==\'Microsoft.Compute/virtualMachines\']" -o json' \
                               .format(tag=tag)
         logger.info('Searching for repair-vm within subscription...')
         output = _call_az_command(find_repair_command)

src/vm-repair/azext_vm_repair/custom.py

@@ -14,6 +14,7 @@
 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 msrestazure.tools import parse_resource_id
+from .exceptions import SkuDoesNotSupportHyperV
 
 from .command_helper_class import command_helper
 from .repair_utils import (
@@ -39,13 +40,15 @@
     _select_distro_linux_gen2,
     _set_repair_map_url,
     _is_gen2,
+    _unlock_encrypted_vm_run,
+    _create_repair_vm,
     _check_n_start_vm
 )
 from .exceptions import AzCommandError, SkuNotAvailableError, UnmanagedDiskCopyError, WindowsOsNotAvailableError, RunScriptNotFoundForIdError, SkuDoesNotSupportHyperV, ScriptReturnsError, SupportingResourceNotFoundError, CommandCanceledByUserError
 logger = get_logger(__name__)
 
 
-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'):
+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', yes=False):
     # Init command helper object
     command = command_helper(logger, cmd, 'vm repair create')
     # Main command calling block
@@ -64,13 +67,12 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
         created_resources = []
 
         # Fetch OS image urn and set OS type for disk create
-        if is_linux:
+        if is_linux and _uses_managed_disk(source_vm):
             # os_image_urn = "UbuntuLTS"
             os_type = 'Linux'
             hyperV_generation_linux = _check_linux_hyperV_gen(source_vm)
             if hyperV_generation_linux == 'V2':
                 logger.info('Generation 2 VM detected, RHEL/Centos/Oracle 6 distros not available to be used for rescue VM ')
-                logger.debug('gen2 machine detected')
                 os_image_urn = _select_distro_linux_gen2(distro)
             else:
                 os_image_urn = _select_distro_linux(distro)
@@ -125,31 +127,31 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
             # Copy OS Disk
             logger.info('Copying OS disk of source VM...')
             copy_disk_id = _call_az_command(copy_disk_command).strip('\n')
-            # For Linux the disk gets not attached at VM creation time. To prevent an incorrect boot state it is required to attach the disk after the VM got created.
-            if not is_linux:
-                # Add copied OS Disk to VM creat command so that the VM is created with the disk attached
-                create_repair_vm_command += ' --attach-data-disks {id}'.format(id=copy_disk_id)
-            # Validate create vm create command to validate parameters before runnning copy disk command
-            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])
-            # Create repair VM
-            logger.info('Creating repair VM...')
-            _call_az_command(create_repair_vm_command, secure_params=[repair_password, repair_username])
 
-            if is_linux:
-                # Attach copied managed disk to new vm
+            # Create VM according to the two conditions: is_linux, unlock_encrypted_vm
+            # Only in the case of a Linux VM without encryption the data-disk gets attached after VM creation.
+            # This is required to prevent an incorrect boot due to an UUID mismatch
+            if not is_linux:
+                # windows
+                _create_repair_vm(copy_disk_id, create_repair_vm_command, repair_password, repair_username)
+
+            if not is_linux and unlock_encrypted_vm:
+                # windows with encryption
+                _create_repair_vm(copy_disk_id, create_repair_vm_command, repair_password, repair_username)
+                _unlock_encrypted_vm_run(repair_vm_name, repair_group_name, is_linux)
+
+            if is_linux and unlock_encrypted_vm:
+                # linux with encryption
+                _create_repair_vm(copy_disk_id, create_repair_vm_command, repair_password, repair_username)
+                _unlock_encrypted_vm_run(repair_vm_name, repair_group_name, is_linux)
+
+            if is_linux and (not unlock_encrypted_vm):
+                # linux without encryption
+                _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...')
                 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)
                 _call_az_command(attach_disk_command)
 
-            # Handle encrypted VM cases
-            if unlock_encrypted_vm:
-                stdout, stderr = _unlock_singlepass_encrypted_disk(repair_vm_name, repair_group_name, is_linux)
-                logger.debug('Unlock script STDOUT:\n%s', stdout)
-                if stderr:
-                    logger.warning('Encryption unlock script error was generated:\n%s', stderr)
-
         # UNMANAGED DISK
         else:
             logger.info('Source VM uses unmanaged disks. Creating repair VM with unmanaged disks.\n')

src/vm-repair/azext_vm_repair/repair_utils.py

@@ -358,7 +358,7 @@ def _check_linux_hyperV_gen(source_vm):
                         .format(i=disk_id)
     hyperVGen = loads(_call_az_command(show_disk_command))
     if hyperVGen != 'V2':
-        logger.info('Trying to check on the source VM if it has the parameter of gen2')
+        logger.info('Checking if source VM is gen2')
         # if image is created from Marketplace gen2 image , the disk will not have the mark for gen2
         fetch_hypervgen_command = 'az vm get-instance-view --ids {id} --query "[instanceView.hyperVGeneration]" -o json'.format(id=source_vm.id)
         hyperVGen_list = loads(_call_az_command(fetch_hypervgen_command))
@@ -682,3 +682,19 @@ def _get_function_param_dict(frame):
         if param in values:
             values[param] = '********'
     return values
+
+
+def _unlock_encrypted_vm_run(repair_vm_name, repair_group_name, is_linux):
+    stdout, stderr = _unlock_singlepass_encrypted_disk(repair_vm_name, repair_group_name, is_linux)
+    logger.debug('Unlock script STDOUT:\n%s', stdout)
+    if stderr:
+        logger.warning('Encryption unlock script error was generated:\n%s', stderr)
+
+
+def _create_repair_vm(copy_disk_id, create_repair_vm_command, repair_password, repair_username, fix_uuid=False):
+    if not fix_uuid:
+        create_repair_vm_command += ' --attach-data-disks {id}'.format(id=copy_disk_id)
+    logger.info('Validating VM template before continuing...')
+    _call_az_command(create_repair_vm_command + ' --validate', secure_params=[repair_password, repair_username])
+    logger.info('Creating repair VM...')
+    _call_az_command(create_repair_vm_command, secure_params=[repair_password, repair_username])