[Compute] Add disk-encryption-set command group

src/azure-cli/HISTORY.rst

@@ -9,7 +9,10 @@ Release History
 
 **Compute**
 
-* vmss create/update: Add --scale-in-policy, which decides which virtual machines are chosen for removal when a VMSS is scaled-in
+* vmss create/update: Add --scale-in-policy, which decides which virtual machines are chosen for removal when a VMSS is scaled-in.
+* Add disk-encryption-set command group.
+* disk create: Add --encryption-type and --disk-encryption-set.
+* vm create: Add --os-disk-encryption-set and --data-disk-encryption-sets.
 
 **Install**
 

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,21 @@
     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.
+"""
+
+helps['disk-encryption-set list'] = """
+type: command
+short-summary: List Disk Encryption Sets.
+"""
+
 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:
@@ -311,6 +314,8 @@ def load_arguments(self, _):
     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.')
         c.argument('attach_data_disks', nargs='+', help='Attach existing data disks to the VM. Can use the name or ID of a managed disk or the URI to an unmanaged disk VHD.')
+        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 of disk encryption sets for data disks.')
 
     with self.argument_context('vm create', arg_group='Dedicated Host', min_api='2019-03-01') as c:
         c.argument('dedicated_host_group', options_list=['--host-group'], is_preview=True, help="Name of the dedicated host group containing the dedicated host this VM will reside in.")
@@ -838,3 +843,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,10 @@ 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 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
 

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) as g:
+        g.custom_command('create', 'create_disk_encryption_set')
+        g.command('delete', 'delete')
+        # g.command('update', 'update')
+        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')

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

@@ -249,8 +249,13 @@ def create_managed_disk(cmd, resource_group_name, disk_name, location=None,  # p
                         # below are generated internally from 'source'
                         source_blob_uri=None, source_disk=None, source_snapshot=None,
                         source_storage_account_id=None, no_wait=False, tags=None, zone=None,
-                        disk_iops_read_write=None, disk_mbps_read_write=None, hyper_v_generation=None):
-    Disk, CreationData, DiskCreateOption = cmd.get_models('Disk', 'CreationData', 'DiskCreateOption')
+                        disk_iops_read_write=None, disk_mbps_read_write=None, hyper_v_generation=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
+
+    Disk, CreationData, DiskCreateOption, Encryption = cmd.get_models(
+        'Disk', 'CreationData', 'DiskCreateOption', 'Encryption')
 
     location = location or _get_resource_group_location(cmd.cli_ctx, resource_group_name)
     if source_blob_uri:
@@ -263,8 +268,6 @@ def create_managed_disk(cmd, resource_group_name, disk_name, location=None,  # p
         option = DiskCreateOption.empty
 
     if source_storage_account_id is None and source_blob_uri is not None:
-        from azure.cli.core.commands.client_factory import get_subscription_id
-        from msrestazure.tools import resource_id
         subscription_id = get_subscription_id(cmd.cli_ctx)
         storage_account_name = source_blob_uri.split('.')[0].split('/')[-1]
         source_storage_account_id = resource_id(
@@ -282,8 +285,21 @@ def create_managed_disk(cmd, resource_group_name, disk_name, location=None,  # p
 
     if size_gb is None and upload_size_bytes is None and (option == DiskCreateOption.empty or for_upload):
         raise CLIError('usage error: --size-gb or --upload-size-bytes required to create an empty disk')
+
+    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
+
     disk = Disk(location=location, creation_data=creation_data, tags=(tags or {}),
-                sku=_get_sku_object(cmd, sku), disk_size_gb=size_gb, os_type=os_type)
+                sku=_get_sku_object(cmd, sku), disk_size_gb=size_gb, os_type=os_type, encryption=encryption)
 
     if hyper_v_generation:
         disk.hyper_vgeneration = hyper_v_generation
@@ -532,7 +548,8 @@ def create_vm(cmd, vm_name, resource_group_name, image=None, size='Standard_DS1_
               identity_role='Contributor', identity_role_id=None, application_security_groups=None, zone=None,
               boot_diagnostics_storage=None, ultra_ssd_enabled=None, ephemeral_os_disk=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):
+              priority=None, max_price=None, eviction_policy=None, enable_agent=None, workspace=None, vmss=None,
+              os_disk_encryption_set=None, data_disk_encryption_sets=None):
     from azure.cli.core.commands.client_factory import get_subscription_id
     from azure.cli.core.util import random_string, hash_string
     from azure.cli.core.commands.arm import ArmTemplateBuilder
@@ -545,9 +562,23 @@ def create_vm(cmd, vm_name, resource_group_name, image=None, size='Standard_DS1_
                                                                 build_vm_daExtensionName_resource)
     from msrestazure.tools import resource_id, is_valid_resource_id
 
+    subscription_id = get_subscription_id(cmd.cli_ctx)
+
+    if os_disk_encryption_set is not None and not is_valid_resource_id(os_disk_encryption_set):
+        os_disk_encryption_set = resource_id(
+            subscription=subscription_id, resource_group=resource_group_name,
+            namespace='Microsoft.Compute', type='diskEncryptionSets', name=os_disk_encryption_set)
+
+    if data_disk_encryption_sets is None:
+        data_disk_encryption_sets = []
+    for i, des in enumerate(data_disk_encryption_sets):
+        if des is not None and not is_valid_resource_id(des):
+            data_disk_encryption_sets[i] = resource_id(
+                subscription=subscription_id, resource_group=resource_group_name,
+                namespace='Microsoft.Compute', type='diskEncryptionSets', name=des)
+
     storage_sku = disk_info['os'].get('storageAccountType')
 
-    subscription_id = get_subscription_id(cmd.cli_ctx)
     network_id_template = resource_id(
         subscription=subscription_id, resource_group=resource_group_name,
         namespace='Microsoft.Network')
@@ -660,7 +691,8 @@ def create_vm(cmd, vm_name, resource_group_name, image=None, size='Standard_DS1_
         boot_diagnostics_storage_uri=boot_diagnostics_storage, ultra_ssd_enabled=ultra_ssd_enabled,
         proximity_placement_group=proximity_placement_group, computer_name=computer_name,
         dedicated_host=dedicated_host, priority=priority, max_price=max_price, eviction_policy=eviction_policy,
-        enable_agent=enable_agent, vmss=vmss)
+        enable_agent=enable_agent, vmss=vmss, os_disk_encryption_set=os_disk_encryption_set,
+        data_disk_encryption_sets=data_disk_encryption_sets)
 
     vm_resource['dependsOn'] = vm_dependencies
 
@@ -2933,3 +2965,28 @@ def execute_query_for_vm(cmd, client, resource_group_name, vm_name, analytics_qu
                        'Please check the status of log analytics workpsace.')
     return client.query(workspace, QueryBody(query=analytics_query, timespan=timespan))
 # endregion
+
+
+# disk encryption set
+def create_disk_encryption_set(cmd, client, resource_group_name, disk_encryption_set_name,
+                               key_url, source_vault, location=None, tags=None):
+    from msrestazure.tools import resource_id, is_valid_resource_id
+    DiskEncryptionSet, EncryptionSetIdentity, KeyVaultAndKeyReference, SourceVault = cmd.get_models(
+        'DiskEncryptionSet', 'EncryptionSetIdentity', 'KeyVaultAndKeyReference', 'SourceVault')
+    encryption_set_identity = EncryptionSetIdentity(type='SystemAssigned')
+    if not is_valid_resource_id(source_vault):
+        source_vault = resource_id(subscription=client.config.subscription_id, resource_group=resource_group_name,
+                                   namespace='Microsoft.KeyVault', type='vaults', name=source_vault)
+    source_vault = SourceVault(id=source_vault)
+    keyVault_and_key_reference = KeyVaultAndKeyReference(source_vault=source_vault, key_url=key_url)
+    disk_encryption_set = DiskEncryptionSet(location=location, tags=tags, identity=encryption_set_identity,
+                                            active_key=keyVault_and_key_reference)
+    return client.create_or_update(resource_group_name, disk_encryption_set_name, disk_encryption_set)
+
+
+def list_disk_encryption_sets(cmd, client, resource_group_name=None):
+    if resource_group_name:
+        return client.list_by_resource_group(resource_group_name)
+    return client.list()
+
+# endregion