v1 tags and public IP fix

src/vm-repair/azext_vm_repair/_help.py

@@ -15,10 +15,17 @@
 helps['vm repair create'] = """
     type: command
     short-summary: Create a new repair VM and attach the source VM's copied OS disk as a data disk.
+    parameters:
+        - name: --tags
+          type: string
+          short-summary: Space-separated tags in 'key[=value]' format. Use '' to clear existing tags.
     examples:
         - name: Create a repair VM
           text: >
             az vm repair create -g MyResourceGroup -n myVM --verbose
+        - name: Create a repair VM with tags
+          text: >
+            az vm repair create -g MyResourceGroup -n myVM --tags env=dev owner=alice --verbose
         - name: Create a repair VM and set the VM authentication
           text: >
             az vm repair create -g MyResourceGroup -n myVM --repair-username username --repair-password password!234 --verbose
@@ -114,15 +121,30 @@
 helps['vm repair repair-and-restore'] = """
     type: command
     short-summary: Repair and restore the VM.
+    parameters:
+        - name: --tags
+          type: string
+          short-summary: Space-separated tags in 'key[=value]' format. Use '' to clear existing tags.
     examples:
+        - name: Repair and restore a VM with tags.
+          text: >
+            az vm repair repair-and-restore --name vmrepairtest --resource-group MyResourceGroup --tags env=prod owner=bob --verbose
         - name: Repair and restore a VM.
           text: >
             az vm repair repair-and-restore --name vmrepairtest --resource-group MyResourceGroup --verbose
 """
+
 helps['vm repair repair-button'] = """
     type: command
     short-summary: repair button script.
+    parameters:
+        - name: --tags
+          type: string
+          short-summary: Space-separated tags in 'key[=value]' format. Use '' to clear existing tags.
     examples:
+        - name: repair-button with tags.
+          text: >
+            az vm repair repair-button --name vmrepairtest --resource-group MyResourceGroup --button-command fstab --tags env=test --verbose
         - name: repair-button.
           text: >
             az vm repair repair-button --name vmrepairtest --resource-group MyResourceGroup --button-command fstab --verbose

src/vm-repair/azext_vm_repair/_params.py

@@ -36,6 +36,7 @@ def load_arguments(self, _):
         c.argument('yes', help='Option to skip prompt for associating public ip in no Tty mode')
         c.argument('disable_trusted_launch', help='Option to disable Trusted Launch security type on the repair vm by setting the security type to Standard.')
         c.argument('os_disk_type', help='Change the OS Disk storage type from the default of PremiumSSD_LRS to the given value.')
+        c.argument('tags', help="Space-separated tags in 'key[=value]' format. Use '' to clear existing tags.")
 
     with self.argument_context('vm repair restore') as c:
         c.argument('repair_vm_id', help='Repair VM resource id.')
@@ -64,6 +65,7 @@ def load_arguments(self, _):
         c.argument('repair_vm_name', help='Name of repair VM.')
         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('tags', help="Space-separated tags in 'key[=value]' format. Use '' to clear existing tags.")
 
     with self.argument_context('vm repair repair-button') as c:
         c.argument('button_command', help='Button_command for repair VM.')
@@ -73,3 +75,4 @@ def load_arguments(self, _):
         c.argument('repair_vm_name', help='Name of repair VM.')
         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('tags', help="Space-separated tags in 'key[=value]' format. Use '' to clear existing tags.")

src/vm-repair/azext_vm_repair/_validators.py

@@ -52,7 +52,7 @@ def validate_create(cmd, namespace):
     else:
         namespace.copy_disk_name = namespace.vm_name + '-DiskCopy-' + timestamp
 
-    # Check copy resouce group name
+    # Check copy resource group name
     if namespace.repair_group_name:
         if namespace.repair_group_name == namespace.resource_group_name:
             raise CLIError('The repair resource group name cannot be the same as the source VM resource group.')
@@ -189,7 +189,7 @@ def validate_reset_nic(cmd, namespace):
             _call_az_command(set_sub_command)
         except AzCommandError as azCommandError:
             logger.error(azCommandError)
-            raise CLIError('Unexpected error occured while setting the subscription..')
+            raise CLIError('Unexpected error occurred while setting the subscription..')
     _validate_and_get_vm(cmd, namespace.resource_group_name, namespace.vm_name)
 
 
@@ -435,4 +435,4 @@ def validate_repair_and_restore(cmd, namespace):
     namespace.associate_public_ip = False
     # 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

@@ -56,7 +56,7 @@
 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="", disable_trusted_launch=False, os_disk_type=None):  
+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, os_disk_type=None, tags=None):  
     """  
     This function creates a repair VM.  
 
@@ -76,7 +76,8 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
     - yes: If True, confirmation prompts will be skipped. Default is False.  
     - 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. 
+    - 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.
     """  
 
     # Logging all the command parameters, except the sensitive data.
@@ -150,12 +151,27 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
         public_ip_name = _make_public_ip_name(repair_vm_name, associate_public_ip)
 
         # Set up base create vm command
+        # Azure CLI accepts tags as either a space-separated list of key=value pairs or a single string. Support both dict and string for flexibility.
+        tag_arg = ''
+        if tags:
+            if isinstance(tags, dict):
+                tag_items = [f"{k}={v}" for k, v in tags.items()]
+                tag_arg = ' --tags ' + ' '.join(tag_items)
+            else:
+                tag_arg = f' --tags {tags}'
+        # Only include the --public-ip-address argument if associate_public_ip is True. Omitting this argument prevents creation of an unwanted public IP.
+        public_ip_arg = f' --public-ip-address {public_ip_name}' if associate_public_ip else ''
         if is_linux:
-            create_repair_vm_command = 'az vm create -g {g} -n {n} --tag {tag} --image {image} --admin-username {username} --admin-password {password} --public-ip-address {option} --custom-data {cloud_init_script}' \
-                .format(g=repair_group_name, n=repair_vm_name, tag=resource_tag, image=os_image_urn, username=repair_username, password=repair_password, option=public_ip_name, cloud_init_script=_get_cloud_init_script())
+            create_repair_vm_command = (
+                f'az vm create -g {repair_group_name} -n {repair_vm_name} --tag {resource_tag} --image {os_image_urn} '
+                f'--admin-username {repair_username} --admin-password {repair_password}{public_ip_arg} '
+                f'--custom-data {_get_cloud_init_script()}{tag_arg}'
+            )
         else:
-            create_repair_vm_command = 'az vm create -g {g} -n {n} --tag {tag} --image {image} --admin-username {username} --admin-password {password} --public-ip-address {option}' \
-                .format(g=repair_group_name, n=repair_vm_name, tag=resource_tag, image=os_image_urn, username=repair_username, password=repair_password, option=public_ip_name)
+            create_repair_vm_command = (
+                f'az vm create -g {repair_group_name} -n {repair_vm_name} --tag {resource_tag} --image {os_image_urn} '
+                f'--admin-username {repair_username} --admin-password {repair_password}{public_ip_arg}{tag_arg}'
+            )
 
         # Fetching the size of the repair VM.  
         sku = _fetch_compatible_sku(source_vm, enable_nested)  
@@ -906,7 +922,7 @@ def reset_nic(cmd, vm_name, resource_group_name, yes=False):
     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):    
+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):    
     """    
     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   
@@ -920,6 +936,7 @@ def repair_and_restore(cmd, vm_name, resource_group_name, repair_password=None,
     :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.  
     """  
     from datetime import datetime  
     import secrets  
@@ -948,7 +965,7 @@ def repair_and_restore(cmd, vm_name, resource_group_name, repair_password=None,
     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, yes=True)  
+    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, yes=True, tags=tags)  
 
     # Log the output of the create operation  
     logger.info('create_out: %s', create_out)  
@@ -965,12 +982,11 @@ def repair_and_restore(cmd, vm_name, resource_group_name, repair_password=None,
         # 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:  
-        # If running the fstab script fails, log the error and clean up resources  
-        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."  
+    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  
@@ -1022,7 +1038,10 @@ def repair_and_restore(cmd, vm_name, resource_group_name, repair_password=None,
     # 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):
+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):
+    """
+    Button-triggered repair operation. Supports tags for the repair VM.
+    """
     from datetime import datetime
     import secrets
     import string

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

@@ -49,6 +49,26 @@ def test_vmrepair_WinManagedCreateRestore(self, resource_group):
         source_vm = vms[0]
         assert source_vm['storageProfile']['osDisk']['name'] == result['copied_disk_name']
 
+        # Test create with tags
+        result = self.cmd('vm repair create -g {rg} -n {vm} --repair-username azureadmin --repair-password !Passw0rd2018 --tags env=test owner=alice -o json --yes').get_output_in_json()
+        assert result['status'] == STATUS_SUCCESS, result['error_message']
+
+        # Check repair VM and tags
+        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]
+        assert repair_vm['storageProfile']['dataDisks'][0]['name'] == result['copied_disk_name']
+        assert repair_vm['tags']['env'] == 'test'
+        assert repair_vm['tags']['owner'] == 'alice'
+
+        # 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.WindowsUnmanaged
 class WindowsUnmanagedDiskCreateRestoreTest(LiveScenarioTest):
 
@@ -118,6 +138,26 @@ def test_vmrepair_LinuxManagedCreateRestore(self, resource_group):
         source_vm = vms[0]
         assert source_vm['storageProfile']['osDisk']['name'] == result['copied_disk_name']
 
+        # Test create with tags
+        result = self.cmd('vm repair create -g {rg} -n {vm} --repair-username azureadmin --repair-password !Passw0rd2018 --tags env=dev owner=bob --yes -o json').get_output_in_json()
+        assert result['status'] == STATUS_SUCCESS, result['error_message']
+
+        # Check repair VM and tags
+        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]
+        assert repair_vm['storageProfile']['dataDisks'][0]['name'] == result['copied_disk_name']
+        assert repair_vm['tags']['env'] == 'dev'
+        assert repair_vm['tags']['owner'] == 'bob'
+
+        # 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.linuxUnmanaged
 class LinuxUnmanagedDiskCreateRestoreTest(LiveScenarioTest):
@@ -399,6 +439,7 @@ def test_vmrepair_LinuxSinglepassKekEncryptedManagedDiskCreateRestore(self, reso
         source_vm = vms[0]
         assert source_vm['storageProfile']['osDisk']['name'] == result['copied_disk_name']
 
+
 @pytest.mark.WindowsNoKekRestore
 class WindowsSinglepassNoKekEncryptedManagedDiskCreateRestoreTest(LiveScenarioTest):
 
@@ -444,6 +485,7 @@ def test_vmrepair_WinSinglepassNoKekEncryptedManagedDiskCreateRestore(self, reso
         source_vm = vms[0]
         assert source_vm['storageProfile']['osDisk']['name'] == result['copied_disk_name']
 
+
 @pytest.mark.LinuxNoKekRestore
 class LinuxSinglepassNoKekEncryptedManagedDiskCreateRestoreTest(LiveScenarioTest):
 
@@ -491,6 +533,7 @@ def test_vmrepair_LinuxSinglepassNoKekEncryptedManagedDiskCreateRestoreTest(self
         source_vm = vms[0]
         assert source_vm['storageProfile']['osDisk']['name'] == result['copied_disk_name']
 
+
 @pytest.mark.WindHelloWorld
 class WindowsRunHelloWorldTest(LiveScenarioTest):
 

src/vm-repair/tests/test_vm_repair_commands.py

@@ -0,0 +1,12 @@
+from azure.cli.testsdk import (ScenarioTest, ResourceGroupPreparer)
+
+class VmRepairTests(ScenarioTest):
+
+    @ResourceGroupPreparer(name_prefix='cli_test_vm_repair')
+    def test_vm_repair_create_success(self, resource_group):
+        self.kwargs.update({
+            'vm_name': 'MyVM'
+        })
+
+        result = self.cmd('az vm repair create -g {rg} -n {vm_name} --verbose').get_output_in_json()
+        self.assertIn('repairVM', result)