diff --git a/src/vm-repair/azext_vm_repair/_help.py b/src/vm-repair/azext_vm_repair/_help.py index 6b39a82fc14..80e69c3169c 100644 --- a/src/vm-repair/azext_vm_repair/_help.py +++ b/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 diff --git a/src/vm-repair/azext_vm_repair/_params.py b/src/vm-repair/azext_vm_repair/_params.py index 5a0caffa0c0..2a57aead1ea 100644 --- a/src/vm-repair/azext_vm_repair/_params.py +++ b/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.") diff --git a/src/vm-repair/azext_vm_repair/_validators.py b/src/vm-repair/azext_vm_repair/_validators.py index d7d42c9f794..cb77fe981b5 100644 --- a/src/vm-repair/azext_vm_repair/_validators.py +++ b/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) \ No newline at end of file diff --git a/src/vm-repair/azext_vm_repair/custom.py b/src/vm-repair/azext_vm_repair/custom.py index 326998ebc0f..b8e8bf7617e 100644 --- a/src/vm-repair/azext_vm_repair/custom.py +++ b/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 diff --git a/src/vm-repair/azext_vm_repair/tests/latest/test_repair_commands.py b/src/vm-repair/azext_vm_repair/tests/latest/test_repair_commands.py index 89c0475d2f5..508dfa8569d 100644 --- a/src/vm-repair/azext_vm_repair/tests/latest/test_repair_commands.py +++ b/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): diff --git a/src/vm-repair/tests/test_vm_repair_commands.py b/src/vm-repair/tests/test_vm_repair_commands.py new file mode 100644 index 00000000000..96d294c1341 --- /dev/null +++ b/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)