vm-repair: improve repair vm creation and encryption unlock [11/11 Hackathon work] (Azure#2659)

* Remove ADE extension for encrypted VM scenario
Attach data disk while vm creation
Fix availability zone bug

* Fix styling

* Suppress linux mount errors

Author: swbae31

src/vm-repair/azext_vm_repair/custom.py

@@ -40,7 +40,6 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
 
     # Init command helper object
     command = command_helper(logger, cmd, 'vm repair create')
-
     # Main command calling block
     try:
         # Fetch source VM data
@@ -69,6 +68,11 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
             raise SkuNotAvailableError('Failed to find compatible VM size for source VM\'s OS disk within given region and subscription.')
         create_repair_vm_command += ' --size {sku}'.format(sku=sku)
 
+        # Set availability zone for vm
+        if source_vm.zones:
+            zone = source_vm.zones[0]
+            create_repair_vm_command += ' --zone {zone}'.format(zone=zone)
+
         # Create new resource group
         create_resource_group_command = 'az group create -l {loc} -n {group_name}' \
                                         .format(loc=source_vm.location, group_name=repair_group_name)
@@ -86,25 +90,29 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
             # Only add hyperV variable when available
             if hyperV_generation:
                 copy_disk_command += ' --hyper-v-generation {hyperV}'.format(hyperV=hyperV_generation)
+            # Set availability zone for vm when available
+            if source_vm.zones:
+                zone = source_vm.zones[0]
+                copy_disk_command += ' --zone {zone}'.format(zone=zone)
+            # Copy OS Disk
+            logger.info('Copying OS disk of source VM...')
+            copy_disk_id = _call_az_command(copy_disk_command).strip('\n')
+            # 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])
-            logger.info('Copying OS disk of source VM...')
-            copy_disk_id = _call_az_command(copy_disk_command).strip('\n')
-
-            attach_disk_command = 'az vm disk attach -g {g} --vm-name {repair} --name {id}' \
-                                  .format(g=repair_group_name, repair=repair_vm_name, id=copy_disk_id)
-
+            # Create repair VM
             logger.info('Creating repair VM...')
             _call_az_command(create_repair_vm_command, secure_params=[repair_password, repair_username])
-            logger.info('Attaching copied disk to repair VM...')
-            _call_az_command(attach_disk_command)
 
             # Handle encrypted VM cases
             if unlock_encrypted_vm:
-                _unlock_singlepass_encrypted_disk(source_vm, resource_group_name, repair_vm_name, repair_group_name, copy_disk_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)
 
         # UNMANAGED DISK
         else:

src/vm-repair/azext_vm_repair/repair_utils.py

@@ -271,7 +271,18 @@ def _secret_tag_check(resource_group_name, copy_disk_name, secreturl):
         _call_az_command(set_tag_command)
 
 
-def _unlock_singlepass_encrypted_disk(source_vm, resource_group_name, repair_vm_name, repair_group_name, copy_disk_name, is_linux):
+def _unlock_singlepass_encrypted_disk(repair_vm_name, repair_group_name, is_linux):
+    logger.info('Unlocking attached copied disk...')
+    if is_linux:
+        return _unlock_mount_linux_encrypted_disk(repair_vm_name, repair_group_name)
+    return _unlock_mount_windows_encrypted_disk(repair_vm_name, repair_group_name)
+
+
+def _unlock_singlepass_encrypted_disk_fallback(source_vm, resource_group_name, repair_vm_name, repair_group_name, copy_disk_name, is_linux):
+    """
+    Fallback for unlocking disk when script fails. This will install the ADE extension to unlock the Data disk.
+    """
+
     # Installs the extension on repair VM and mounts the disk after unlocking.
     encryption_type, key_vault, kekurl, secreturl = _fetch_encryption_settings(source_vm)
     if is_linux:
@@ -296,25 +307,31 @@ def _unlock_singlepass_encrypted_disk(source_vm, resource_group_name, repair_vm_
             # Validating secret tag and setting original tag if it got changed
             _secret_tag_check(resource_group_name, copy_disk_name, secreturl)
             logger.debug("Manually unlocking and mounting disk for Linux VMs.")
-            _manually_unlock_mount_encrypted_disk(repair_vm_name, repair_group_name)
+            _unlock_mount_linux_encrypted_disk(repair_vm_name, repair_group_name)
     except AzCommandError as azCommandError:
         error_message = str(azCommandError)
         # Linux VM encryption extension bug where it fails and then continue to mount disk manually
         if is_linux and "Failed to encrypt data volumes with error" in error_message:
             logger.debug("Expected bug for linux VMs. Ignoring error.")
             # Validating secret tag and setting original tag if it got changed
             _secret_tag_check(resource_group_name, copy_disk_name, secreturl)
-            _manually_unlock_mount_encrypted_disk(repair_vm_name, repair_group_name)
+            _unlock_mount_linux_encrypted_disk(repair_vm_name, repair_group_name)
         else:
             raise
 
 
-def _manually_unlock_mount_encrypted_disk(repair_vm_name, repair_group_name):
+def _unlock_mount_linux_encrypted_disk(repair_vm_name, repair_group_name):
     # Unlocks the disk using the phasephrase and mounts it on the repair VM.
-    LINUX_RUN_SCRIPT_NAME = 'mount-encrypted-disk.sh'
+    LINUX_RUN_SCRIPT_NAME = 'linux-mount-encrypted-disk.sh'
     return _invoke_run_command(LINUX_RUN_SCRIPT_NAME, repair_vm_name, repair_group_name, True)
 
 
+def _unlock_mount_windows_encrypted_disk(repair_vm_name, repair_group_name):
+    # Unlocks the disk using the phasephrase and mounts it on the repair VM.
+    WINDOWS_RUN_SCRIPT_NAME = 'win-mount-encrypted-disk.ps1'
+    return _invoke_run_command(WINDOWS_RUN_SCRIPT_NAME, repair_vm_name, repair_group_name, False)
+
+
 def _fetch_compatible_windows_os_urn(source_vm):
     location = source_vm.location
     fetch_urn_command = 'az vm image list -s "2016-Datacenter" -f WindowsServer -p MicrosoftWindowsServer -l {loc} --verbose --all --query "[?sku==\'2016-Datacenter\'].urn | reverse(sort(@))" -o json'.format(loc=location)