[Backup] az backup restore restore-disks: Add support for enabling Disk access settings for managed VM restores

src/azure-cli/azure/cli/command_modules/backup/_params.py

@@ -33,6 +33,7 @@
 allowed_immutability_options = ['Disabled', 'Locked', 'Unlocked']
 enable_disable_options = ['Enable', 'Disable']
 enable_disable_permadisable_options = ['Enable', 'Disable', 'PermanentlyDisable']
+allowed_disk_access_options = ['EnablePrivateAccessForAllDisks', 'EnablePublicAccessForAllDisks', 'SameAsOnSourceDisks']
 
 backup_management_type_help = """Specify the backup management type. Define how Azure Backup manages the backup of entities within the ARM resource. For eg: AzureWorkloads refers to workloads installed within Azure VMs, AzureStorage refers to entities within Storage account. Required only if friendly name is used as Container name."""
 container_name_help = """Name of the backup container. Accepts 'Name' or 'FriendlyName' from the output of az backup container list command. If 'FriendlyName' is passed then BackupManagementType is required."""
@@ -399,6 +400,8 @@ def load_arguments(self, _):
         c.argument('storage_account_resource_group', help='Name of the resource group which contains the storage account. Default value will be same as --resource-group if not specified.')
         c.argument('restore_to_edge_zone', arg_type=get_three_state_flag(), help='Switch parameter to indicate edge zone VM restore. This parameter can\'t be used in cross region and cross subscription restore scenarios.')
         c.argument('tenant_id', help='ID of the tenant if the Resource Guard protecting the vault exists in a different tenant.')
+        c.argument('disk_access_option', arg_type=get_enum_type(allowed_disk_access_options), help='Specify the disk access option for target disks.')
+        c.argument('target_disk_access_id', help='Specify the target disk access ID when --disk-access-option is set to EnablePrivateAccessForAllDisks')
 
     with self.argument_context('backup restore restore-azurefileshare') as c:
         c.argument('resolve_conflict', resolve_conflict_type)

src/azure-cli/azure/cli/command_modules/backup/custom.py

@@ -27,7 +27,7 @@
     BackupResourceVaultConfig, BackupResourceVaultConfigResource, DiskExclusionProperties, ExtendedProperties, \
     MoveRPAcrossTiersRequest, RecoveryPointRehydrationInfo, IaasVMRestoreWithRehydrationRequest, IdentityInfo, \
     BackupStatusRequest, ListRecoveryPointsRecommendedForMoveRequest, IdentityBasedRestoreDetails, ScheduleRunType, \
-    UnlockDeleteRequest, ResourceGuardProxyBase, ResourceGuardProxyBaseResource
+    UnlockDeleteRequest, ResourceGuardProxyBase, ResourceGuardProxyBaseResource, TargetDiskNetworkAccessSettings
 from azure.mgmt.recoveryservicesbackup.passivestamp.models import CrrJobRequest, CrossRegionRestoreRequest
 
 import azure.cli.command_modules.backup._validators as validators
@@ -1340,6 +1340,44 @@ def _get_alr_restore_mode(target_vm_name, target_vnet_name, target_vnet_resource
         """)
 
 
+def _set_pe_restore_trigger_restore_properties(cmd, trigger_restore_properties, disk_access_option, target_disk_access_id,
+                                               recovery_point, use_secondary_region):
+    if not hasattr(recovery_point.properties, 'is_private_access_enabled_on_any_disk'):
+        return trigger_restore_properties
+    if recovery_point.properties.is_private_access_enabled_on_any_disk:
+        if disk_access_option is None:
+            raise InvalidArgumentValueError("--disk-access-option parameter must be provided since private access "
+                                            "is enabled in given recovery point")
+
+        if disk_access_option == "EnablePrivateAccessForAllDisks":
+            if target_disk_access_id is None:
+                raise InvalidArgumentValueError("--target-disk-access-id must be provided when --disk-access-option "
+                                                "is set to EnablePrivateAccessForAllDisks")
+
+        if disk_access_option == "SameAsOnSourceDisks":
+            if use_secondary_region:
+                raise InvalidArgumentValueError("Given --disk-access-option is not applicable to cross region restore")
+            if target_disk_access_id is not None:
+                raise InvalidArgumentValueError("--target-disk-access-id can't be provided for the "
+                                                "given --disk-access-option")
+
+        if disk_access_option == "EnablePublicAccessForAllDisks":
+            if target_disk_access_id is not None:
+                raise InvalidArgumentValueError("--target-disk-access-id can't be provided for the "
+                                                "given --disk-access-option")
+
+        trigger_restore_properties.target_disk_network_access_settings = TargetDiskNetworkAccessSettings(
+            target_disk_access_id=target_disk_access_id,
+            target_disk_network_access_option=disk_access_option
+        )
+    else:
+        if disk_access_option is not None or target_disk_access_id is not None:
+            raise InvalidArgumentValueError("--disk-access-option parameter can't be provided since private access "
+                                            "is not enabled in given recovery point")
+
+    return trigger_restore_properties
+
+
 def _set_edge_zones_trigger_restore_properties(cmd, trigger_restore_properties, restore_to_edge_zone, recovery_point,
                                                target_subscription, use_secondary_region, restore_mode):
     # TODO: As the subscription we currently use does not have access to Edge Zones, no tests have been written for
@@ -1375,7 +1413,7 @@ def restore_disks(cmd, client, resource_group_name, vault_name, container_name,
                   mi_user_assigned=None, target_zone=None, restore_mode='AlternateLocation', target_vm_name=None,
                   target_vnet_name=None, target_vnet_resource_group=None, target_subnet_name=None,
                   target_subscription_id=None, storage_account_resource_group=None, restore_to_edge_zone=None,
-                  tenant_id=None):
+                  tenant_id=None, disk_access_option=None, target_disk_access_id=None):
     vault = vaults_cf(cmd.cli_ctx).get(resource_group_name, vault_name)
     vault_location = vault.location
     vault_identity = vault.identity
@@ -1481,6 +1519,10 @@ def restore_disks(cmd, client, resource_group_name, vault_name, container_name,
                                                                             recovery_point, target_subscription,
                                                                             use_secondary_region, restore_mode)
 
+    trigger_restore_properties = _set_pe_restore_trigger_restore_properties(cmd, trigger_restore_properties,
+                                                                            disk_access_option, target_disk_access_id,
+                                                                            recovery_point, use_secondary_region)
+
     trigger_restore_request = RestoreRequestResource(properties=trigger_restore_properties)
 
     if use_secondary_region:

src/azure-cli/azure/cli/command_modules/backup/custom_base.py

@@ -421,7 +421,7 @@ def restore_disks(cmd, client, resource_group_name, vault_name, container_name,
                   mi_user_assigned=None, target_zone=None, restore_mode='AlternateLocation', target_vm_name=None,
                   target_vnet_name=None, target_vnet_resource_group=None, target_subnet_name=None,
                   target_subscription_id=None, storage_account_resource_group=None, restore_to_edge_zone=None,
-                  tenant_id=None):
+                  tenant_id=None, disk_access_option=None, target_disk_access_id=None):
 
     if rehydration_duration < 10 or rehydration_duration > 30:
         raise InvalidArgumentValueError('--rehydration-duration must have a value between 10 and 30 (both inclusive).')
@@ -439,7 +439,7 @@ def restore_disks(cmd, client, resource_group_name, vault_name, container_name,
                                 mi_system_assigned, mi_user_assigned, target_zone, restore_mode, target_vm_name,
                                 target_vnet_name, target_vnet_resource_group, target_subnet_name,
                                 target_subscription_id, storage_account_resource_group, restore_to_edge_zone,
-                                tenant_id)
+                                tenant_id, disk_access_option, target_disk_access_id)
 
 
 def enable_for_azurefileshare(cmd, client, resource_group_name, vault_name, policy_name, storage_account,