[Compute] Add disk-encryption-set command group

src/azure-cli/HISTORY.rst

@@ -18,6 +18,9 @@ Release History
 * vmss create/update: Add --scale-in-policy, which decides which virtual machines are chosen for removal when a VMSS is scaled-in.
 * vm/vmss update: Add --priority.
 * vm/vmss update: Add --max-price.
+* Add disk-encryption-set command group (create, show, update, delete, list).
+* disk create: Add --encryption-type and --disk-encryption-set.
+* vm/vmss create: Add --os-disk-encryption-set and --data-disk-encryption-sets.
 
 **Core**
 

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

@@ -135,3 +135,7 @@ def cf_log_analytics_data_plane(cli_ctx, _):
     cred, _, _ = profile.get_login_credentials(
         resource="https://api.loganalytics.io")
     return LogAnalyticsDataClient(cred)
+
+
+def cf_disk_encryption_set(cli_ctx, _):
+    return _compute_client_factory(cli_ctx).disk_encryption_sets

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

@@ -100,6 +100,39 @@
     crafted: true
 """
 
+helps['disk-encryption-set'] = """
+type: group
+short-summary: Disk Encryption Set resource.
+"""
+
+helps['disk-encryption-set create'] = """
+type: command
+short-summary: Create a disk encryption set.
+examples:
+  - name: Create a disk encryption set
+    text: az disk-encryption-set create -g MyResourceGroup -n MyDiskEncryptionSet --key-url MyKey --source-vault MyVault
+"""
+
+helps['disk-encryption-set delete'] = """
+type: command
+short-summary: Delete a disk encryption set.
+"""
+
+helps['disk-encryption-set list'] = """
+type: command
+short-summary: List disk encryption sets.
+"""
+
+helps['disk-encryption-set show'] = """
+type: command
+short-summary: Get information of a disk encryption sets.
+"""
+
+helps['disk-encryption-set update'] = """
+type: command
+short-summary: Update a disk encryption set.
+"""
+
 helps['image'] = """
 type: group
 short-summary: Manage custom virtual machine images.

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

@@ -54,6 +54,7 @@ def load_arguments(self, _):
 
     extension_instance_name_type = CLIArgumentType(help="Name of extension instance, which can be customized. Default: name of the extension.")
     image_template_name_type = CLIArgumentType(overrides=name_arg_type, id_part='name')
+    disk_encryption_set_name = CLIArgumentType(overrides=name_arg_type, help='Name of disk encryption set.', id_part='name')
 
     # StorageAccountTypes renamed to DiskStorageAccountTypes in 2018_06_01 of azure-mgmt-compute
     DiskStorageAccountTypes = DiskStorageAccountTypes or StorageAccountTypes
@@ -109,6 +110,8 @@ 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.')
 
     for scope in ['disk create', 'snapshot create']:
         with self.argument_context(scope) as c:
@@ -660,6 +663,9 @@ def load_arguments(self, _):
             c.argument('ultra_ssd_enabled', ultra_ssd_enabled_type)
             c.argument('ephemeral_os_disk', arg_type=get_three_state_flag(), min_api='2018-06-01',
                        help='Allows you to create an OS disk directly on the host node, providing local disk performance and faster VM/VMSS reimage time.', is_preview=True)
+            c.argument('os_disk_encryption_set', help='Name or ID of disk encryption set for OS disk.')
+            c.argument('data_disk_encryption_sets', nargs='+',
+                       help='Names or IDs (space delimited) of disk encryption sets for data disks.')
 
         with self.argument_context(scope, arg_group='Network') as c:
             c.argument('vnet_name', help='Name of the virtual network when creating a new one or referencing an existing one.')
@@ -839,3 +845,12 @@ def load_arguments(self, _):
         c.argument('analytics_query', options_list=['--analytics-query', '-q'], help="Query to execute over Log Analytics data.")
         c.argument('timespan', help="Timespan over which to query. Defaults to querying all available data.")
     # endregion
+
+    # region disk encryption set
+    with self.argument_context('disk-encryption-set') as c:
+        c.argument('disk_encryption_set_name', disk_encryption_set_name)
+        c.argument('key_url', help='URL pointing to a key or secret in KeyVault.')
+        c.argument('source_vault', help='Name or ID of the KeyVault containing the key or secret.')
+        c.argument('location', validator=get_default_location_from_resource_group)
+        c.argument('tags', tags_type)
+    # endregion

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

@@ -6,6 +6,8 @@
 
 from enum import Enum
 
+from knack.util import CLIError
+
 from azure.cli.core.util import b64encode
 from azure.cli.core.profiles import ResourceType
 from azure.cli.core.commands.arm import ArmTemplateBuilder
@@ -249,7 +251,7 @@ def build_vm_resource(  # pylint: disable=too-many-locals, too-many-statements
         attach_os_disk=None, os_disk_size_gb=None, custom_data=None, secrets=None, license_type=None, zone=None,
         disk_info=None, boot_diagnostics_storage_uri=None, ultra_ssd_enabled=None, proximity_placement_group=None,
         computer_name=None, dedicated_host=None, priority=None, max_price=None, eviction_policy=None,
-        enable_agent=None, vmss=None):
+        enable_agent=None, vmss=None, os_disk_encryption_set=None, data_disk_encryption_sets=None):
 
     os_caching = disk_info['os'].get('caching')
 
@@ -337,7 +339,12 @@ def _build_storage_profile():
                     'createOption': 'fromImage',
                     'name': os_disk_name,
                     'caching': os_caching,
-                    'managedDisk': {'storageAccountType': disk_info['os'].get('storageAccountType')}
+                    'managedDisk': {
+                        'storageAccountType': disk_info['os'].get('storageAccountType'),
+                        'diskEncryptionSet': {
+                            'id': os_disk_encryption_set,
+                        }
+                    }
                 },
                 'imageReference': {
                     'publisher': os_publisher,
@@ -351,7 +358,12 @@ def _build_storage_profile():
                     'createOption': 'fromImage',
                     'name': os_disk_name,
                     'caching': os_caching,
-                    'managedDisk': {'storageAccountType': disk_info['os'].get('storageAccountType')}
+                    'managedDisk': {
+                        'storageAccountType': disk_info['os'].get('storageAccountType'),
+                        'diskEncryptionSet': {
+                            'id': os_disk_encryption_set,
+                        }
+                    }
                 },
                 "imageReference": {
                     'id': image_reference
@@ -373,6 +385,12 @@ def _build_storage_profile():
         if disk_info['os'].get('writeAcceleratorEnabled') is not None:
             profile['osDisk']['writeAcceleratorEnabled'] = disk_info['os']['writeAcceleratorEnabled']
         data_disks = [v for k, v in disk_info.items() if k != 'os']
+        if data_disk_encryption_sets:
+            if len(data_disk_encryption_sets) != len(data_disks):
+                raise CLIError(
+                    'usage error: Number of --data-disk-encryption-sets mismatches with number of data disks.')
+            for i, data_disk in enumerate(data_disks):
+                data_disk['managedDisk']['diskEncryptionSet'] = {'id': data_disk_encryption_sets[i]}
         if data_disks:
             profile['dataDisks'] = data_disks
 
@@ -651,7 +669,8 @@ def build_vmss_resource(cmd, name, naming_prefix, location, tags, overprovision,
                         single_placement_group=None, platform_fault_domain_count=None, custom_data=None,
                         secrets=None, license_type=None, zones=None, priority=None, eviction_policy=None,
                         application_security_groups=None, ultra_ssd_enabled=None, proximity_placement_group=None,
-                        terminate_notification_time=None, max_price=None, scale_in_policy=None):
+                        terminate_notification_time=None, max_price=None, scale_in_policy=None,
+                        os_disk_encryption_set=None, data_disk_encryption_sets=None):
 
     # Build IP configuration
     ip_configuration = {
@@ -715,6 +734,10 @@ def build_vmss_resource(cmd, name, naming_prefix, location, tags, overprovision,
             'caching': os_caching,
             'managedDisk': {'storageAccountType': disk_info['os'].get('storageAccountType')}
         }
+        if os_disk_encryption_set is not None:
+            storage_properties['osDisk']['managedDisk']['diskEncryptionSet'] = {
+                'id': os_disk_encryption_set
+            }
         if disk_info['os'].get('diffDiskSettings'):
             storage_properties['osDisk']['diffDiskSettings'] = disk_info['os']['diffDiskSettings']
 
@@ -730,6 +753,12 @@ def build_vmss_resource(cmd, name, naming_prefix, location, tags, overprovision,
             'id': image
         }
     data_disks = [v for k, v in disk_info.items() if k != 'os']
+    if data_disk_encryption_sets:
+        if len(data_disk_encryption_sets) != len(data_disks):
+            raise CLIError(
+                'usage error: Number of --data-disk-encryption-sets mismatches with number of data disks.')
+        for i, data_disk in enumerate(data_disks):
+            data_disk['managedDisk']['diskEncryptionSet'] = {'id': data_disk_encryption_sets[i]}
     if data_disks:
         storage_properties['dataDisks'] = data_disks
 

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

@@ -13,7 +13,8 @@
                                                           cf_gallery_images, cf_gallery_image_versions,
                                                           cf_proximity_placement_groups,
                                                           cf_dedicated_hosts, cf_dedicated_host_groups,
-                                                          cf_log_analytics_data_plane)
+                                                          cf_log_analytics_data_plane,
+                                                          cf_disk_encryption_set)
 from azure.cli.command_modules.vm._format import (
     transform_ip_addresses, transform_vm, transform_vm_create_output, transform_vm_usage_list, transform_vm_list,
     transform_sku_for_table_output, transform_disk_show_table_output, transform_extension_show_table_output,
@@ -172,6 +173,11 @@ def load_command_table(self, _):
         client_factory=cf_log_analytics_data_plane,
     )
 
+    compute_disk_encryption_set_sdk = CliCommandType(
+        operations_tmpl='azure.mgmt.compute.operations#DiskEncryptionSetsOperations.{}',
+        client_factory=cf_disk_encryption_set
+    )
+
     with self.command_group('disk', compute_disk_sdk, operation_group='disks', min_api='2017-03-30') as g:
         g.custom_command('create', 'create_managed_disk', supports_no_wait=True, table_transformer=transform_disk_show_table_output, validator=process_disk_or_snapshot_create_namespace)
         g.command('delete', 'delete', supports_no_wait=True, confirmation=True)
@@ -182,6 +188,13 @@ 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:
+        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')
+        g.show_command('show', 'get')
+        g.custom_command('list', 'list_disk_encryption_sets')
+
     with self.command_group('image', compute_image_sdk, min_api='2016-04-30-preview') as g:
         g.custom_command('create', 'create_image', validator=process_image_create_namespace)
         g.custom_command('list', 'list_images')