[Compute] Add Managed Identity Support in Azure Disk Encryption

src/azure-cli/azure/cli/command_modules/vm/_help.py

@@ -1602,6 +1602,9 @@
   - name: Create a Debian11 VM with both system and user assigned identity.
     text: >
         az vm create -n MyVm -g rg1 --image Debian11 --assign-identity [system] /subscriptions/99999999-1bf0-4dda-aec3-cb9272f09590/resourcegroups/myRG/providers/Microsoft.ManagedIdentity/userAssignedIdentities/myID
+  - name: Create a vm with user assigned identity and add encryption identity for Azure disk encryption
+    text: >
+        az vm create -n MyVm -g rg1 --image Debian11 --assign-identity myID --encryption-identity /subscriptions/00000000-0000-0000-0000-000000000000/resourcegroups/myRG/providers/Microsoft.ManagedIdentity/userAssignedIdentities/myID
   - name: Create a VM in an availability zone in the current resource group's region.
     supported-profiles: latest
     text: >
@@ -1810,6 +1813,9 @@
   - name: Enable disk encryption on the OS disk and/or data disks. Encrypt mounted disks. (autogenerated)
     text: |
         az vm encryption enable --disk-encryption-keyvault MyVault --name MyVm --resource-group MyResourceGroup --volume-type DATA
+  - name: Add support for using managed identity to authenticate to customer's keyvault for ADE operation
+    text: >
+        az vm encryption enable --disk-encryption-keyvault MyVault --name MyVm --resource-group MyResourceGroup --encryption-identity EncryptionIdentity
     crafted: true
 """
 

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

@@ -485,6 +485,7 @@ def load_arguments(self, _):
         c.argument('enable_vtpm', enable_vtpm_type)
         c.argument('user_data', help='UserData for the VM. It can be passed in as file or string.', completer=FilesCompleter(), type=file_type, min_api='2021-03-01')
         c.argument('enable_hibernation', arg_type=get_three_state_flag(), min_api='2021-03-01', help='The flag that enable or disable hibernation capability on the VM.')
+        c.argument('encryption_identity', help='Resource Id of the user managed identity which can be used for Azure disk encryption')
 
     with self.argument_context('vm create', arg_group='Storage') as c:
         c.argument('attach_os_disk', help='Attach an existing OS disk to the VM. Can use the name or ID of a managed disk or the URI to an unmanaged disk VHD.')
@@ -1194,6 +1195,9 @@ def load_arguments(self, _):
             c.argument('key_encryption_key', help='Key vault key name or URL used to encrypt the disk encryption key.')
             c.argument('key_encryption_keyvault', help='Name or ID of the key vault containing the key encryption key used to encrypt the disk encryption key. If missing, CLI will use `--disk-encryption-keyvault`.')
 
+    with self.argument_context('vm encryption enable') as c:
+        c.argument('encryption_identity', help='Resource Id of the user managed identity which can be used for Azure disk encryption')
+
     for scope in ['vm extension', 'vmss extension']:
         with self.argument_context(scope) as c:
             c.argument('publisher', help='The name of the extension publisher.')

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

@@ -886,8 +886,8 @@ def create_vm(cmd, vm_name, resource_group_name, image=None, size='Standard_DS1_
               storage_account_type=None, vnet_type=None, nsg_type=None, public_ip_address_type=None, nic_type=None,
               validate=False, custom_data=None, secrets=None, plan_name=None, plan_product=None, plan_publisher=None,
               plan_promotion_code=None, license_type=None, assign_identity=None, identity_scope=None,
-              identity_role=None, identity_role_id=None, application_security_groups=None, zone=None,
-              boot_diagnostics_storage=None, ultra_ssd_enabled=None,
+              identity_role=None, identity_role_id=None, encryption_identity=None,
+              application_security_groups=None, zone=None, boot_diagnostics_storage=None, ultra_ssd_enabled=None,
               ephemeral_os_disk=None, ephemeral_os_disk_placement=None,
               proximity_placement_group=None, dedicated_host=None, dedicated_host_group=None, aux_subscriptions=None,
               priority=None, max_price=None, eviction_policy=None, enable_agent=None, workspace=None, vmss=None,
@@ -1145,6 +1145,30 @@ def create_vm(cmd, vm_name, resource_group_name, image=None, size='Standard_DS1_
             master_template.add_resource(build_msi_role_assignment(vm_name, vm_id, identity_role_id,
                                                                    role_assignment_guid, identity_scope))
 
+    if encryption_identity:
+        if not cmd.supported_api_version(min_api='2023-09-01', resource_type=ResourceType.MGMT_COMPUTE):
+            raise CLIError("Usage error: Encryption Identity required API version 2023-09-01 or higher."
+                           "You can set the cloud's profile to use the required API Version with:"
+                           "az cloud set --profile latest --name <cloud name>")
+
+        if 'identity' in vm_resource and 'userAssignedIdentities' in vm_resource['identity'] \
+            and encryption_identity.lower() in \
+                (k.lower() for k in vm_resource['identity']['userAssignedIdentities'].keys()):
+            if 'securityProfile' not in vm_resource['properties']:
+                vm_resource['properties']['securityProfile'] = {}
+            if 'encryptionIdentity' not in vm_resource['properties']['securityProfile']:
+                vm_resource['properties']['securityProfile']['encryptionIdentity'] = {}
+
+            vm_securityProfile_EncryptionIdentity = vm_resource['properties']['securityProfile']['encryptionIdentity']
+
+            if 'userAssignedIdentityResourceId' not in vm_securityProfile_EncryptionIdentity or \
+                    vm_securityProfile_EncryptionIdentity['userAssignedIdentityResourceId'] != encryption_identity:
+                vm_resource['properties']['securityProfile']['encryptionIdentity']['userAssignedIdentityResourceId'] \
+                    = encryption_identity
+        else:
+            raise CLIError("Encryption Identity should be an ARM Resource ID of one of the "
+                           "user assigned identities associated to the resource")
+
     if workspace is not None:
         workspace_id = _prepare_workspace(cmd, resource_group_name, workspace)
         master_template.add_secure_parameter('workspaceId', workspace_id)

src/azure-cli/azure/cli/command_modules/vm/disk_encryption.py

@@ -61,6 +61,47 @@ def _detect_ade_status(vm):
     return True, False   # we believe impossible to have both old & new ADE
 
 
+def updateVmEncryptionSetting(cmd, vm, resource_group_name, vm_name, encryption_identity):
+    from azure.cli.core.azclierror import ArgumentUsageError
+    if vm.identity is None or vm.identity.user_assigned_identities is None or encryption_identity.lower() not in \
+            (k.lower() for k in vm.identity.user_assigned_identities.keys()):
+        raise ArgumentUsageError("Encryption Identity should be an ARM Resource ID of one of the "
+                                 "user assigned identities associated to the resource")
+
+    SecurityProfile, EncryptionIdentity = cmd.get_models('SecurityProfile', 'EncryptionIdentity')
+    updateVm = False
+
+    if vm.security_profile is None:
+        vm.security_profile = SecurityProfile()
+    if vm.security_profile.encryption_identity is None:
+        vm.security_profile.encryption_identity = EncryptionIdentity()
+    if vm.security_profile.encryption_identity.user_assigned_identity_resource_id is None \
+            or vm.security_profile.encryption_identity.user_assigned_identity_resource_id.lower() \
+            != encryption_identity:
+        vm.security_profile.encryption_identity.user_assigned_identity_resource_id = encryption_identity
+        updateVm = True
+
+    if updateVm:
+        compute_client = _compute_client_factory(cmd.cli_ctx)
+        updateEncryptionIdentity \
+            = compute_client.virtual_machines.begin_create_or_update(resource_group_name, vm_name, vm)
+        LongRunningOperation(cmd.cli_ctx)(updateEncryptionIdentity)
+        result = updateEncryptionIdentity.result()
+        return result is not None and result.provisioning_state == 'Succeeded'
+    logger.info("No changes in identity")
+    return True
+
+
+def isVersionSuppprtedForEncryptionIdentity(cmd):
+    from azure.cli.core.profiles import ResourceType
+    from knack.util import CLIError
+    if not cmd.supported_api_version(min_api='2023-09-01', resource_type=ResourceType.MGMT_COMPUTE):
+        raise CLIError("Usage error: Encryption Identity required API version 2023-09-01 or higher."
+                       "You can set the cloud's profile to use the required API Version with:"
+                       "az cloud set --profile latest --name <cloud name>")
+    return True
+
+
 def encrypt_vm(cmd, resource_group_name, vm_name,  # pylint: disable=too-many-locals, too-many-statements
                disk_encryption_keyvault,
                aad_client_id=None,
@@ -70,7 +111,7 @@ def encrypt_vm(cmd, resource_group_name, vm_name,  # pylint: disable=too-many-lo
                key_encryption_algorithm='RSA-OAEP',
                volume_type=None,
                encrypt_format_all=False,
-               force=False):
+               force=False, encryption_identity=None):
     from azure.mgmt.core.tools import parse_resource_id
     from knack.util import CLIError
 
@@ -109,6 +150,12 @@ def encrypt_vm(cmd, resource_group_name, vm_name,  # pylint: disable=too-many-lo
     # disk encryption key itself can be further protected, so let us verify
     if key_encryption_key:
         key_encryption_keyvault = key_encryption_keyvault or disk_encryption_keyvault
+    if encryption_identity and isVersionSuppprtedForEncryptionIdentity(cmd):
+        result = updateVmEncryptionSetting(cmd, vm, resource_group_name, vm_name, encryption_identity)
+        if result:
+            logger.info("Encryption Identity successfully set in virtual machine")
+        else:
+            raise CLIError("Failed to update encryption Identity to the VM")
 
     #  to avoid bad server errors, ensure the vault has the right configurations
     _verify_keyvault_good_for_encryption(cmd.cli_ctx, disk_encryption_keyvault, key_encryption_keyvault, vm, force)
@@ -553,10 +600,19 @@ def _report_client_side_validation_error(msg):
     disk_vault_resource_info = parse_resource_id(disk_vault_id)
     key_vault = client.get(disk_vault_resource_info['resource_group'], disk_vault_resource_info['name'])
 
-    # ensure vault has 'EnabledForDiskEncryption' permission
-    if not key_vault.properties or not key_vault.properties.enabled_for_disk_encryption:
-        _report_client_side_validation_error("Keyvault '{}' is not enabled for disk encryption.".format(
-            disk_vault_resource_info['resource_name']))
+    # ensure vault has 'EnabledForDiskEncryption' permission or VM has encryption identity set for ADE operation
+    if resource_type == 'VM':
+        vm_encryption_identity = vm_or_vmss
+    else:
+        vm_encryption_identity = vm_or_vmss.virtual_machine_profile
+
+    if vm_encryption_identity and vm_encryption_identity.security_profile and \
+        vm_encryption_identity.security_profile.encryption_identity and \
+            vm_encryption_identity.security_profile.encryption_identity.user_assigned_identity_resource_id:
+        pass
+    elif not key_vault.properties or not key_vault.properties.enabled_for_disk_encryption:
+        _report_client_side_validation_error(
+            "Keyvault '{}' is not enabled for disk encryption.".format(disk_vault_resource_info['resource_name']))
 
     if kek_vault_id:
         kek_vault_info = parse_resource_id(kek_vault_id)