dev and test

src/vm-repair/azext_vm_repair/_params.py

@@ -34,6 +34,7 @@ def load_arguments(self, _):
         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)')
         c.argument('yes', help='Option to skip prompt for associating public ip and confirm yes to it in no Tty mode')
+        c.argument('disable_trusted_launch', help='Option to disable Trusted Launch security type on the repair vm .')
 
     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/custom.py

@@ -54,10 +54,10 @@
 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, encrypt_recovery_key=""):
+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="", disable_trusted_launch=False):
 
     # 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)
+    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, disable_trusted_launch: %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, disable_trusted_launch)
 
     # Init command helper object
     command = command_helper(logger, cmd, 'vm repair create')
@@ -118,10 +118,14 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
 
         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
+            if not disable_trusted_launch:
+                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
+            else:
+                logger.debug('Set security-type to Standard...')
+                create_repair_vm_command += ' --security-type Standard'
 
             logger.debug('Fetching OS Disk security profile...')
             osdisk_security_params = _fetch_osdisk_security_profile_parameters(source_vm)

src/vm-repair/azext_vm_repair/repair_utils.py

@@ -416,6 +416,7 @@ def _unlock_singlepass_encrypted_disk(repair_vm_name, repair_group_name, is_linu
 
 def _unlock_singlepass_encrypted_disk_fallback(source_vm, resource_group_name, repair_vm_name, repair_group_name, copy_disk_name, is_linux):
     """
+    This method is not actually invoked. 
     Fallback for unlocking disk when script fails. This will install the ADE extension to unlock the Data disk.
     """
 
@@ -507,13 +508,13 @@ def _fetch_compatible_windows_os_urn_v2(source_vm):
         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...')
+        logger.info('Fetching compatible Windows OS images from gallery V2...')
         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...')
+        logger.info('Fetching compatible Windows OS images from gallery for 2022 Datacenter V2...')
         urns = loads(_call_az_command(fetch_urn_command))
 
     # No OS images available for Windows2016

src/vm-repair/azext_vm_repair/tests/latest/test_repair_commands.py

@@ -844,6 +844,45 @@ def test_vmrepair_ConfidentialVMAndUnlockDisk(self, resource_group):
         # Call Restore
         self.cmd('vm repair restore -g {rg} -n {vm} --yes')
 
+        # Check swapped OS disk
+        vms = self.cmd('vm list -g {rg} -o json').get_output_in_json()
+        source_vm = vms[0]
+        assert source_vm['storageProfile']['osDisk']['name'] == result['copied_disk_name']
+
+
+@pytest.mark.winTrustedLaunchVMFlag
+class WindowsTrustedLaunchVMFlag(LiveScenarioTest):
+
+    @ResourceGroupPreparer(location='westus2')
+    def test_vmrepair_TrustedLaunchVMFlag(self, resource_group):
+        self.kwargs.update({
+            'vm': 'vm1',  
+            'rg': resource_group  
+        })
+
+        # Create test VM
+        # need to create a TL VM
+        self.cmd('vm create -g {rg} -n {vm} --admin-username azureadmin --admin-password !Passw0rd2024 --image Win2022Datacenter')
+        vms = self.cmd('vm list -g {rg} -o json').get_output_in_json()
+        # Something wrong with vm create command if it fails here
+        assert len(vms) == 1
+
+        # Test create
+        result = self.cmd('vm repair create -g {rg} -n {vm} --repair-username azureadmin --repair-password !Passw0rd2024 --yes  --unlock-encrypted-vm --encrypt-recovery-key !Passw0rd2024 --disable-trusted-launch --verbose -o json').get_output_in_json()
+        assert result['status'] == STATUS_SUCCESS, result['error_message']
+
+        # Check repair VM
+        repair_vms = self.cmd('vm list -g {} -o json'.format(result['repair_resource_group'])).get_output_in_json()
+        assert len(repair_vms) == 1
+        repair_vm = repair_vms[0]
+        print("Printing repair VM Output")
+        print(repair_vm)
+        # Check attached data disk
+        assert repair_vm['storageProfile']['dataDisks'][0]['name'] == result['copied_disk_name']
+
+        # Call Restore
+        self.cmd('vm repair restore -g {rg} -n {vm} --yes')
+
         # Check swapped OS disk
         vms = self.cmd('vm list -g {rg} -o json').get_output_in_json()
         source_vm = vms[0]