[Compute] disk update: Add --disk-encryption-set and --encryption-type; snapshot create/update: Add --disk-encryption-set and --encryption-type (#11805)

* [Compute] disk update: Add --disk-encryption-set and --encryption-type

* snapshot create/update: Add --disk-encryption-set and --encryption-type

* Remove preview label

* Fix style problem

* Add parameter validation

* Fix yaml

* Fix test

Co-authored-by: Xiaojian Xu <arrownj@126.com>

Author: qwordy

src/azure-cli/HISTORY.rst

@@ -3,6 +3,11 @@
 Release History
 ===============
 
+**Compute**
+
+* disk update: Add --disk-encryption-set and --encryption-type
+* snapshot create/update: Add --disk-encryption-set and --encryption-type
+
 **Stoarge**
 
 * Upgrade azure-mgmt-storage version to 7.1.0

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

@@ -114,8 +114,9 @@ def load_arguments(self, _):
                 c.argument('hyper_v_generation', arg_type=hyper_v_gen_sku, help='The hypervisor generation of the Virtual Machine. Applicable to OS disks only.')
             else:
                 c.ignore('access_level', 'for_upload', 'hyper_v_generation')
-            c.argument('encryption_type', arg_type=get_enum_type(self.get_models('EncryptionType')), help='Encryption type.')
-            c.argument('disk_encryption_set', help='Name or ID of disk encryption set that is used to encrypt the disk.')
+            c.argument('encryption_type', min_api='2019-07-01', arg_type=get_enum_type(self.get_models('EncryptionType')),
+                       help='Encryption type. EncryptionAtRestWithPlatformKey: Disk is encrypted with XStore managed key at rest. It is the default encryption type. EncryptionAtRestWithCustomerKey: Disk is encrypted with Customer managed key at rest.')
+            c.argument('disk_encryption_set', min_api='2019-07-01', help='Name or ID of disk encryption set that is used to encrypt the disk.')
 
     for scope in ['disk create', 'snapshot create']:
         with self.argument_context(scope) as c:

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

@@ -202,7 +202,7 @@ def load_command_table(self, _):
         g.generic_update_command('update', custom_func_name='update_managed_disk', setter_arg_name='disk', supports_no_wait=True)
         g.wait_command('wait')
 
-    with self.command_group('disk-encryption-set', compute_disk_encryption_set_sdk, client_factory=cf_disk_encryption_set, min_api='2019-07-01', is_preview=True) as g:
+    with self.command_group('disk-encryption-set', compute_disk_encryption_set_sdk, client_factory=cf_disk_encryption_set, min_api='2019-07-01') as g:
         g.custom_command('create', 'create_disk_encryption_set', supports_no_wait=True)
         g.command('delete', 'delete')
         g.generic_update_command('update', custom_func_name='update_disk_encryption_set', setter_arg_name='disk_encryption_set')

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

@@ -327,7 +327,11 @@ def list_managed_disks(cmd, resource_group_name=None):
     return client.disks.list()
 
 
-def update_managed_disk(cmd, instance, size_gb=None, sku=None, disk_iops_read_write=None, disk_mbps_read_write=None):
+def update_managed_disk(cmd, resource_group_name, instance, size_gb=None, sku=None, disk_iops_read_write=None,
+                        disk_mbps_read_write=None, encryption_type=None, disk_encryption_set=None):
+    from msrestazure.tools import resource_id, is_valid_resource_id
+    from azure.cli.core.commands.client_factory import get_subscription_id
+
     if size_gb is not None:
         instance.disk_size_gb = size_gb
     if sku is not None:
@@ -336,6 +340,17 @@ def update_managed_disk(cmd, instance, size_gb=None, sku=None, disk_iops_read_wr
         instance.disk_iops_read_write = disk_iops_read_write
     if disk_mbps_read_write is not None:
         instance.disk_mbps_read_write = disk_mbps_read_write
+    if disk_encryption_set is not None:
+        if instance.encryption.type != 'EncryptionAtRestWithCustomerKey' and \
+                encryption_type != 'EncryptionAtRestWithCustomerKey':
+            raise CLIError('usage error: Please set --encryption-type to EncryptionAtRestWithCustomerKey')
+        if not is_valid_resource_id(disk_encryption_set):
+            disk_encryption_set = resource_id(
+                subscription=get_subscription_id(cmd.cli_ctx), resource_group=resource_group_name,
+                namespace='Microsoft.Compute', type='diskEncryptionSets', name=disk_encryption_set)
+        instance.encryption.disk_encryption_set_id = disk_encryption_set
+    if encryption_type is not None:
+        instance.encryption.type = encryption_type
     return instance
 # endregion
 
@@ -408,12 +423,18 @@ def list_images(cmd, resource_group_name=None):
 
 
 # region Snapshots
+# pylint: disable=unused-argument,too-many-locals
 def create_snapshot(cmd, resource_group_name, snapshot_name, location=None, size_gb=None, sku='Standard_LRS',
-                    source=None, for_upload=None, incremental=None,  # pylint: disable=unused-argument
+                    source=None, for_upload=None, incremental=None,
                     # below are generated internally from 'source'
                     source_blob_uri=None, source_disk=None, source_snapshot=None, source_storage_account_id=None,
-                    hyper_v_generation=None, tags=None, no_wait=False):
-    Snapshot, CreationData, DiskCreateOption = cmd.get_models('Snapshot', 'CreationData', 'DiskCreateOption')
+                    hyper_v_generation=None, tags=None, no_wait=False, disk_encryption_set=None,
+                    encryption_type=None):
+    from msrestazure.tools import resource_id, is_valid_resource_id
+    from azure.cli.core.commands.client_factory import get_subscription_id
+
+    Snapshot, CreationData, DiskCreateOption, Encryption = cmd.get_models(
+        'Snapshot', 'CreationData', 'DiskCreateOption', 'Encryption')
 
     location = location or _get_resource_group_location(cmd.cli_ctx, resource_group_name)
     if source_blob_uri:
@@ -432,8 +453,22 @@ def create_snapshot(cmd, resource_group_name, snapshot_name, location=None, size
 
     if size_gb is None and option == DiskCreateOption.empty:
         raise CLIError('Please supply size for the snapshots')
+
+    if disk_encryption_set is not None and not is_valid_resource_id(disk_encryption_set):
+        disk_encryption_set = resource_id(
+            subscription=get_subscription_id(cmd.cli_ctx), resource_group=resource_group_name,
+            namespace='Microsoft.Compute', type='diskEncryptionSets', name=disk_encryption_set)
+
+    if disk_encryption_set is not None and encryption_type is None:
+        raise CLIError('usage error: Please specify --encryption-type.')
+    if encryption_type is not None:
+        encryption = Encryption(type=encryption_type, disk_encryption_set_id=disk_encryption_set)
+    else:
+        encryption = None
+
     snapshot = Snapshot(location=location, creation_data=creation_data, tags=(tags or {}),
-                        sku=_get_sku_object(cmd, sku), disk_size_gb=size_gb, incremental=incremental)
+                        sku=_get_sku_object(cmd, sku), disk_size_gb=size_gb, incremental=incremental,
+                        encryption=encryption)
     if hyper_v_generation:
         snapshot.hyper_vgeneration = hyper_v_generation
 
@@ -453,9 +488,23 @@ def list_snapshots(cmd, resource_group_name=None):
     return client.snapshots.list()
 
 
-def update_snapshot(cmd, instance, sku=None):
+def update_snapshot(cmd, resource_group_name, instance, sku=None, disk_encryption_set=None, encryption_type=None):
+    from msrestazure.tools import resource_id, is_valid_resource_id
+    from azure.cli.core.commands.client_factory import get_subscription_id
+
     if sku is not None:
         _set_sku(cmd, instance, sku)
+    if disk_encryption_set is not None:
+        if instance.encryption.type != 'EncryptionAtRestWithCustomerKey' and \
+                encryption_type != 'EncryptionAtRestWithCustomerKey':
+            raise CLIError('usage error: Please set --encryption-type to EncryptionAtRestWithCustomerKey')
+        if not is_valid_resource_id(disk_encryption_set):
+            disk_encryption_set = resource_id(
+                subscription=get_subscription_id(cmd.cli_ctx), resource_group=resource_group_name,
+                namespace='Microsoft.Compute', type='diskEncryptionSets', name=disk_encryption_set)
+        instance.encryption.disk_encryption_set_id = disk_encryption_set
+    if encryption_type is not None:
+        instance.encryption.type = encryption_type
     return instance
 # endregion