[Compute] Add Managed Identity Support in Azure Disk Encryption for VMSS

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

@@ -2643,6 +2643,9 @@
   - name: Create a Debian11 VM scaleset with a user assigned identity.
     text: >
         az vmss create -n MyVmss -g rg1 --image Debian11 --assign-identity  /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 vmss 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 --orchestration-mode Uniform --lb-sku Standard
   - name: Create a Debian11 VM scaleset with both system and user assigned identity.
     text: >
         az vmss create -n MyVmss -g rg1 --image Debian11 --assign-identity  [system] /subscriptions/99999999-1bf0-4dda-aec3-cb9272f09590/resourcegroups/myRG/providers/Microsoft.ManagedIdentity/userAssignedIdentities/myID
@@ -2760,6 +2763,9 @@
   - name: encrypt a VM scale set using a key vault in the same resource group
     text: >
         az vmss encryption enable -g MyResourceGroup -n MyVmss --disk-encryption-keyvault MyVault
+  - name: Add support for using managed identity to authenticate to customer's keyvault for ADE operation
+    text: >
+        az vmss encryption enable --disk-encryption-keyvault MyVault --name MyVm --resource-group MyResourceGroup --encryption-identity EncryptionIdentity
   - name: Encrypt a VMSS with managed disks. (autogenerated)
     text: |
         az vmss encryption enable --disk-encryption-keyvault MyVault --name MyVmss --resource-group MyResourceGroup --volume-type DATA

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

@@ -1191,8 +1191,13 @@ 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 create', 'vmss create']:
+        with self.argument_context(scope) 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 encryption enable', 'vmss encryption enable']:
+        with self.argument_context(scope) 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:

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

@@ -3171,7 +3171,7 @@ def create_vmss(cmd, vmss_name, resource_group_name, image=None,
                 public_ip_address_type=None, storage_profile=None,
                 single_placement_group=None, custom_data=None, secrets=None, platform_fault_domain_count=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,
+                assign_identity=None, identity_scope=None, identity_role=None, encryption_identity=None,
                 identity_role_id=None, zones=None, priority=None, eviction_policy=None,
                 application_security_groups=None, ultra_ssd_enabled=None,
                 ephemeral_os_disk=None, ephemeral_os_disk_placement=None,
@@ -3532,6 +3532,35 @@ def _get_public_ip_address_allocation(value, sku):
                 role_assignment_guid = str(_gen_guid())
                 master_template.add_resource(build_msi_role_assignment(vmss_name, vmss_id, identity_role_id,
                                                                        role_assignment_guid, identity_scope, False))
+        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 vmss_resource and 'userAssignedIdentities' in vmss_resource['identity'] \
+                and encryption_identity.lower() in \
+                    (k.lower() for k in vmss_resource['identity']['userAssignedIdentities'].keys()):
+
+                if 'virtualMachineProfile' not in vmss_resource['properties']:
+                    vmss_resource['properties']['virtualMachineProfile'] = {}
+                if 'securityProfile' not in vmss_resource['properties']['virtualMachineProfile']:
+                    vmss_resource['properties']['virtualMachineProfile']['securityProfile'] = {}
+                if 'encryptionIdentity' not in vmss_resource['properties']['virtualMachineProfile']['securityProfile']:
+                    vmss_resource['properties']['virtualMachineProfile']['securityProfile']['encryptionIdentity'] = {}
+
+                vmss_securityProfile_EncryptionIdentity \
+                    = vmss_resource['properties']['virtualMachineProfile']['securityProfile']['encryptionIdentity']
+
+                if 'userAssignedIdentityResourceId' not in vmss_securityProfile_EncryptionIdentity or \
+                        vmss_securityProfile_EncryptionIdentity['userAssignedIdentityResourceId'] \
+                        != encryption_identity:
+                    vmss_securityProfile_EncryptionIdentity['userAssignedIdentityResourceId'] = encryption_identity
+                    vmss_resource['properties']['virtualMachineProfile']['securityProfile']['encryptionIdentity'] \
+                        = vmss_securityProfile_EncryptionIdentity
+            else:
+                raise CLIError("Encryption Identity should be an ARM Resource ID of one of the "
+                               "user assigned identities associated to the resource")
     else:
         raise CLIError('usage error: --orchestration-mode (Uniform | Flexible)')
 

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

@@ -92,6 +92,41 @@ def updateVmEncryptionSetting(cmd, vm, resource_group_name, vm_name, encryption_
     return True
 
 
+def updateVmssEncryptionSetting(cmd, vmss, resource_group_name, vmss_name, encryption_identity):
+    from azure.cli.core.azclierror import ArgumentUsageError
+    if vmss.identity is None or vmss.identity.user_assigned_identities is None or encryption_identity.lower() not in \
+            (k.lower() for k in vmss.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")
+
+    VirtualMachineProfile, SecurityProfile, EncryptionIdentity \
+        = cmd.get_models('VirtualMachineProfile', 'SecurityProfile', 'EncryptionIdentity')
+    updateVmss = False
+
+    if vmss.virtual_machine_profile is None:
+        vmss.virtual_machine_profile = VirtualMachineProfile()
+    if vmss.virtual_machine_profile.security_profile is None:
+        vmss.virtual_machine_profile.security_profile = SecurityProfile()
+    if vmss.virtual_machine_profile.security_profile.encryption_identity is None:
+        vmss.virtual_machine_profile.security_profile.encryption_identity = EncryptionIdentity()
+    if vmss.virtual_machine_profile.security_profile.encryption_identity.user_assigned_identity_resource_id\
+        is None or vmss.virtual_machine_profile.security_profile.encryption_identity\
+            .user_assigned_identity_resource_id.lower() != encryption_identity:
+        vmss.virtual_machine_profile.security_profile.encryption_identity.user_assigned_identity_resource_id \
+            = encryption_identity
+        updateVmss = True
+
+    if updateVmss:
+        compute_client = _compute_client_factory(cmd.cli_ctx)
+        updateEncryptionIdentity \
+            = compute_client.virtual_machine_scale_sets.begin_create_or_update(resource_group_name, vmss_name, vmss)
+        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
@@ -436,8 +471,10 @@ def encrypt_vmss(cmd, resource_group_name, vmss_name,  # pylint: disable=too-man
                  key_encryption_key=None,
                  key_encryption_algorithm='RSA-OAEP',
                  volume_type=None,
+                 encryption_identity=None,
                  force=False):
     from azure.mgmt.core.tools import parse_resource_id
+    from knack.util import CLIError
 
     # pylint: disable=no-member
     UpgradeMode, VirtualMachineScaleSetExtension, VirtualMachineScaleSetExtensionProfile = cmd.get_models(
@@ -459,6 +496,13 @@ def encrypt_vmss(cmd, resource_group_name, vmss_name,  # pylint: disable=too-man
     if key_encryption_key:
         key_encryption_keyvault = key_encryption_keyvault or disk_encryption_keyvault
 
+    if encryption_identity and isVersionSuppprtedForEncryptionIdentity(cmd):
+        result = updateVmssEncryptionSetting(cmd, vmss, resource_group_name, vmss_name, encryption_identity)
+        if result:
+            logger.info("Encryption Identity successfully set in virtual machine scale set")
+        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, vmss, force)