Add CVM functionality to vm repair: unlocking disks (#7839)

* copy from dataa PR

* change param name to be shorter per check fail

* resolving linter issues

* try to resolve linter errors

* fixing style attempt

* resolve test errors and fix regression

* fix test

* changed error type

* changelog

* fixing error message

* setup version

* fixed changelog syntax

* Update src/vm-repair/HISTORY.rst

Co-authored-by: Xing Zhou <Zhou.Xing@microsoft.com>

* changelog edit

---------

Co-authored-by: Xing Zhou <Zhou.Xing@microsoft.com>

src/vm-repair/HISTORY.rst

@@ -6,6 +6,7 @@ Release History
 ++++++
 Fixed and updated several vm-repair tests for better coverage. 
 Removed and updated broken image aliases pointing at images that no longer existed. 
+Add `--encrypt-recovery-key` string parameter to `vm repair create` to use recovery key provided by the user to unlock the disk for a confidential VM. 
 
 1.0.8
 ++++++
@@ -19,7 +20,6 @@ az command adjustment
 ++++++
 Add CLI update wait for ASG to wait for the operation done as the async 2rd operation will cancel the 1st call.
 
-
 1.0.5
 ++++++
 Bug fix ASG is not added properly when reset the nic

src/vm-repair/azext_vm_repair/_params.py

@@ -29,6 +29,7 @@ def load_arguments(self, _):
         c.argument('copy_disk_name', help='Name of OS disk copy.')
         c.argument('repair_group_name', help='Name for new or existing resource group that will contain repair VM.')
         c.argument('unlock_encrypted_vm', help='Option to auto-unlock encrypted VMs using current subscription auth.')
+        c.argument('encrypt_recovery_key', help='Option to auto-unlock encrypted VMs using provided recovery password.')
         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|sles12|sles15|ubuntu20|centos7|centos8|oracle7)')

src/vm-repair/azext_vm_repair/_validators.py

@@ -8,7 +8,7 @@
 from re import match, search, findall
 from knack.log import get_logger
 from knack.util import CLIError
-from azure.cli.core.azclierror import ValidationError
+from azure.cli.core.azclierror import ValidationError, RequiredArgumentMissingError
 
 from azure.cli.command_modules.vm.custom import get_vm, _is_linux_os
 from azure.cli.command_modules.resource._client_factory import _resource_client_factory
@@ -72,6 +72,10 @@ def validate_create(cmd, namespace):
     else:
         logger.debug('The source VM\'s OS disk is not encrypted')
 
+    if namespace.encrypt_recovery_key:
+        if not namespace.unlock_encrypted_vm:
+            raise RequiredArgumentMissingError('Recovery password is provided in the argument, but --unlock-encrypted-vm is not passed. Rerun command adding --unlock-encrypted-vm.')
+
     if namespace.enable_nested:
         if is_linux:
             raise CLIError('Nested VM is not supported for Linux VM')
@@ -429,4 +433,4 @@ def validate_repair_and_restore(cmd, namespace):
 
     # Validate repair run command
     source_vm = _validate_and_get_vm(cmd, namespace.resource_group_name, namespace.vm_name)
-    is_linux = _is_linux_os(source_vm)
+    is_linux = _is_linux_os(source_vm)

src/vm-repair/azext_vm_repair/custom.py

@@ -45,15 +45,18 @@
     _check_n_start_vm,
     _check_existing_rg,
     _fetch_architecture,
-    _select_distro_linux_Arm64
+    _select_distro_linux_Arm64,
+    _fetch_vm_security_profile_parameters,
+    _fetch_osdisk_security_profile_parameters,
+    _fetch_compatible_windows_os_urn_v2
 )
 from .exceptions import AzCommandError, RunScriptNotFoundForIdError, 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', yes=False):
+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, encrypt_recovery_key=""):
 
-    # log all the parameters
+    # log all the parameters not logging the bitlocker key
     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, yes: %s', vm_name, resource_group_name, repair_password, repair_username, repair_vm_name, copy_disk_name, repair_group_name, unlock_encrypted_vm, enable_nested, associate_public_ip, distro, yes)
 
     # Init command helper object
@@ -88,7 +91,10 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
             else:
                 os_image_urn = _select_distro_linux(distro)
         else:
-            os_image_urn = _fetch_compatible_windows_os_urn(source_vm)
+            if encrypt_recovery_key:
+                os_image_urn = _fetch_compatible_windows_os_urn_v2(source_vm)
+            else:
+                os_image_urn = _fetch_compatible_windows_os_urn(source_vm)
             os_type = 'Windows'
 
         # Set up base create vm command
@@ -110,6 +116,18 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
             zone = source_vm.zones[0]
             create_repair_vm_command += ' --zone {zone}'.format(zone=zone)
 
+        if encrypt_recovery_key:
+            # For confidential VM and Trusted VM security tags some of the SKU expects the right security type, secure_boot_enabled and vtpm_enabled
+            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
+
+            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
+
         # Create new resource group
         existing_rg = _check_existing_rg(repair_group_name)
         if not existing_rg:
@@ -151,7 +169,7 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
             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)
+                _unlock_encrypted_vm_run(repair_vm_name, repair_group_name, is_linux, encrypt_recovery_key)
 
             if is_linux and unlock_encrypted_vm:
                 # linux with encryption
@@ -612,18 +630,17 @@ def reset_nic(cmd, vm_name, resource_group_name, yes=False):
         vnet_resource_group = subnet_id_tokens[-7]
         ipconfig_name = ip_config_object['name']
         orig_ip_address = ip_config_object['privateIPAddress']
-        application_names=""
-        applicationSecurityGroups='applicationSecurityGroups'
+        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]+ " "
-                    
+                    application_names += application_id_tokens[-1] + " "
+
         logger.info('applicationSecurityGroups {application_names}...\n')
 
-
         # Dynamic | Static
         orig_ip_allocation_method = ip_config_object['privateIPAllocationMethod']
 
@@ -640,19 +657,18 @@ def reset_nic(cmd, vm_name, resource_group_name, yes=False):
         # 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)
+                                .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 updated
+
+        # Wait for IP 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)
+            .format(g=resource_group_name, nic=primary_nic_name)
         _call_az_command(wait_ip_update_command)
 
-
         # 4) Change things back. This will also invoke and wait for a VM restart.
         logger.info('NIC reset is complete. Now reverting back to your original configuration...\n')
         # If user had dynamic config, change back to dynamic
@@ -661,15 +677,15 @@ def reset_nic(cmd, vm_name, resource_group_name, yes=False):
             # Revert Static to Dynamic
             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)
+                                    .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:
             # Revert to original static ip
             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)
+                                    .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)

src/vm-repair/azext_vm_repair/repair_utils.py

@@ -407,11 +407,11 @@ def _secret_tag_check(resource_group_name, copy_disk_name, secreturl):
         _call_az_command(set_tag_command)
 
 
-def _unlock_singlepass_encrypted_disk(repair_vm_name, repair_group_name, is_linux):
+def _unlock_singlepass_encrypted_disk(repair_vm_name, repair_group_name, is_linux, encrypt_recovery_key):
     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)
+    return _unlock_mount_windows_encrypted_disk(repair_vm_name, repair_group_name, encrypt_recovery_key)
 
 
 def _unlock_singlepass_encrypted_disk_fallback(source_vm, resource_group_name, repair_vm_name, repair_group_name, copy_disk_name, is_linux):
@@ -462,8 +462,15 @@ def _unlock_mount_linux_encrypted_disk(repair_vm_name, repair_group_name):
     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):
+def _unlock_mount_windows_encrypted_disk(repair_vm_name, repair_group_name, encrypt_recovery_key):
     # Unlocks the disk using the phasephrase and mounts it on the repair VM.
+    if encrypt_recovery_key:
+        logger.info('Using bitlocker password to unlock...')
+        WINDOWS_RUN_SCRIPT_NAME = 'win-mount-encrypted-disk-bitlockerV.ps1'
+        BITLOCKER_RECOVERY_PARAMS = []
+        BITLOCKER_RECOVERY_PARAMS.append('bitlockerkey="{}"'.format(encrypt_recovery_key))
+        return _invoke_run_command(WINDOWS_RUN_SCRIPT_NAME, repair_vm_name, repair_group_name, False, parameters=BITLOCKER_RECOVERY_PARAMS)
+
     WINDOWS_RUN_SCRIPT_NAME = 'win-mount-encrypted-disk.ps1'
     return _invoke_run_command(WINDOWS_RUN_SCRIPT_NAME, repair_vm_name, repair_group_name, False)
 
@@ -491,6 +498,32 @@ def _fetch_compatible_windows_os_urn(source_vm):
     return urns[0]
 
 
+def _fetch_compatible_windows_os_urn_v2(source_vm):
+    location = source_vm.location
+
+    # We will prefer to fetch image using source vm sku, that we match the CVM requirements.
+    if source_vm.storage_profile is not None and source_vm.storage_profile.image_reference is not None:
+        sku = source_vm.storage_profile.image_reference.sku
+        offer = source_vm.storage_profile.image_reference.offer
+        publisher = source_vm.storage_profile.image_reference.publisher
+        fetch_urn_command = 'az vm image list -s {sku} -f {offer} -p {publisher} -l {loc} --verbose --all --query "[?sku==\'{sku}\'].urn | reverse(sort(@))" -o json'.format(loc=location, sku=sku, offer=offer, publisher=publisher)
+        logger.info('Fetching compatible Windows OS images from gallery...')
+        urns = loads(_call_az_command(fetch_urn_command))
+
+    if not urns or len(urns) == 0:
+        # If source SKU not available then defaulting 2022 datacenter image.
+        fetch_urn_command = 'az vm image list -s "2022-Datacenter" -f WindowsServer -p MicrosoftWindowsServer -l {loc} --verbose --all --query "[?sku==\'2022-datacenter\'].urn | reverse(sort(@))" -o json'.format(loc=location)
+        logger.info('Fetching compatible Windows OS images from gallery for 2022 Datacenter...')
+        urns = loads(_call_az_command(fetch_urn_command))
+
+    # No OS images available for Windows2016
+    if not urns:
+        raise WindowsOsNotAvailableError()
+    logger.debug('Fetched Urns:\n%s', urns)
+    logger.debug('Defaulting to first image available. Returning Urn 0: %s', urns[0])
+    return urns[0]
+
+
 def _select_distro_linux(distro):
     image_lookup = {
         'rhel7': 'RedHat:rhel-raw:7-raw:latest',
@@ -682,18 +715,19 @@ def _get_function_param_dict(frame):
     _, _, _, values = inspect.getargvalues(frame)
     if 'cmd' in values:
         del values['cmd']
-    secure_params = ['repair_password', 'repair_username']
+    secure_params = ['repair_password', 'repair_username', 'encrypt_recovery_key']
     for param in secure_params:
         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)
+def _unlock_encrypted_vm_run(repair_vm_name, repair_group_name, is_linux, encrypt_recovery_key = ""):
+    stdout, stderr = _unlock_singlepass_encrypted_disk(repair_vm_name, repair_group_name, is_linux, encrypt_recovery_key)
     logger.debug('Unlock script STDOUT:\n%s', stdout)
     if stderr:
         logger.warning('Encryption unlock script error was generated:\n%s', stderr)
+        raise Exception('Unexpected error occured while unlocking encrypted disk.')
 
 
 def _create_repair_vm(copy_disk_id, create_repair_vm_command, repair_password, repair_username, fix_uuid=False):
@@ -726,3 +760,40 @@ def _fetch_architecture(source_vm):
     architecture = loads(_call_az_command(architecture_type_cmd).strip('\n'))
 
     return architecture[0][0]
+
+
+def _fetch_non_standard_security_type(source_vm):
+    """
+    Returns security type if security type is not standard and needs to be set.
+    """
+    if source_vm.security_profile is None or source_vm.security_profile.security_type is None:
+        return
+    if source_vm.security_profile.security_type.lower() == "standard":
+        return
+    return source_vm.security_profile.security_type
+
+
+def _fetch_vm_security_profile_parameters(source_vm):
+    create_repair_vm_command = ''
+    non_standard_security_type = _fetch_non_standard_security_type(source_vm)
+    if non_standard_security_type is None:
+        return create_repair_vm_command
+    create_repair_vm_command += ' --security-type {securityType}'.format(securityType=non_standard_security_type)
+    if source_vm.security_profile.uefi_settings is not None:
+        if source_vm.security_profile.uefi_settings.secure_boot_enabled is not None:
+            create_repair_vm_command += ' --enable-secure-boot {enableSecureBoot}'.format(enableSecureBoot=source_vm.security_profile.uefi_settings.secure_boot_enabled)
+
+        if source_vm.security_profile.uefi_settings.v_tpm_enabled is not None:
+            create_repair_vm_command += ' --enable-vtpm {enableVTpm}'.format(enableVTpm=source_vm.security_profile.uefi_settings.v_tpm_enabled)
+    return create_repair_vm_command
+
+
+def _fetch_osdisk_security_profile_parameters(source_vm):
+    create_repair_vm_command = ''
+    if source_vm.storage_profile.os_disk.managed_disk is not None and source_vm.storage_profile.os_disk.managed_disk.security_profile is not None:
+        create_repair_vm_command += ' --os-disk-security-encryption-type {val}'.format(val=source_vm.storage_profile.os_disk.managed_disk.security_profile.security_encryption_type)
+
+        if source_vm.storage_profile.os_disk.managed_disk.security_profile.disk_encryption_set is not None:
+            create_repair_vm_command += ' --os-disk-secure-vm-disk-encryption-set {val}'.format(val=source_vm.storage_profile.os_disk.managed_disk.security_profile.disk_encryption_set.id)
+
+    return create_repair_vm_command