Support creating cosmos db containers with client encryption policy

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

@@ -182,6 +182,7 @@ def load_arguments(self, _):
         c.argument('partition_key_version', type=int, options_list=['--partition-key-version'], help='The version of partition key.')
         c.argument('default_ttl', options_list=['--ttl'], type=int, help='Default TTL. If the value is missing or set to "-1", items don’t expire. If the value is set to "n", items will expire "n" seconds after last modified time.')
         c.argument('indexing_policy', options_list=['--idx'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Indexing Policy, you can enter it as a string or as a file, e.g., --idx @policy-file.json or ' + SQL_GREMLIN_INDEXING_POLICY_EXAMPLE)
+        c.argument('client_encryption_policy', options_list=['--client-encryption-policy'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Client Encryption Policy, you can enter it as a string or as a file, e.g., --client-encryption-policy @policy-file.json. ')
         c.argument('unique_key_policy', options_list=['--unique-key-policy', '-u'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Unique Key Policy, you can enter it as a string or as a file, e.g., --unique-key-policy @policy-file.json or ' + SQL_UNIQUE_KEY_POLICY_EXAMPLE)
         c.argument('conflict_resolution_policy', options_list=['--conflict-resolution-policy', '-c'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Conflict Resolution Policy, you can enter it as a string or as a file, e.g., --conflict-resolution-policy @policy-file.json or ' + SQL_GREMLIN_CONFLICT_RESOLUTION_POLICY_EXAMPLE)
         c.argument('max_throughput', max_throughput_type)

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

@@ -10,6 +10,7 @@
 from knack.util import CLIError
 from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
 from azure.cli.core.util import sdk_no_wait
+import json 
 
 from azure.mgmt.cosmosdb.models import (
     ConsistencyPolicy,
@@ -24,6 +25,8 @@
     SqlContainerResource,
     SqlContainerCreateUpdateParameters,
     ContainerPartitionKey,
+    ClientEncryptionIncludedPath,
+    ClientEncryptionPolicy,
     ResourceIdentityType,
     SqlStoredProcedureResource,
     SqlStoredProcedureCreateUpdateParameters,
@@ -516,11 +519,12 @@ def _populate_sql_container_definition(sql_container_resource,
                                        default_ttl,
                                        indexing_policy,
                                        unique_key_policy,
+                                       client_encryption_policy,
                                        partition_key_version,
                                        conflict_resolution_policy,
                                        analytical_storage_ttl):
     if all(arg is None for arg in
-           [partition_key_path, partition_key_version, default_ttl, indexing_policy, unique_key_policy, conflict_resolution_policy, analytical_storage_ttl]):
+           [partition_key_path, partition_key_version, default_ttl, indexing_policy, unique_key_policy, client_encryption_policy, conflict_resolution_policy, analytical_storage_ttl]):
         return False
 
     if partition_key_path is not None:
@@ -538,7 +542,11 @@ def _populate_sql_container_definition(sql_container_resource,
         sql_container_resource.indexing_policy = indexing_policy
 
     if unique_key_policy is not None:
-        sql_container_resource.unique_key_policy = unique_key_policy
+        sql_container_resource.unique_key_policy = unique_key_policy    
+
+    if client_encryption_policy is not None:
+        _validate_and_populate_client_encryption_policy(sql_container_resource,
+                                                    client_encryption_policy)
 
     if conflict_resolution_policy is not None:
         sql_container_resource.conflict_resolution_policy = conflict_resolution_policy
@@ -548,6 +556,93 @@ def _populate_sql_container_definition(sql_container_resource,
 
     return True
 
+def _validate_and_populate_client_encryption_policy(sql_container_resource,
+                                                    client_encryption_policy):
+    if client_encryption_policy is not None:
+        from azure.cli.core.util import get_file_json, shell_safe_json_parse
+        data = shell_safe_json_parse(json.dumps(client_encryption_policy))        
+
+        if "includedPaths" in data:
+                includedpaths = data['includedPaths']
+        else:
+            raise CLIError(None, f"includedPaths missing in Client Encryption Policy. Please verify your Client Encryption Policy JSON string")
+
+        if "policyFormatVersion" in data:
+                policyFormatVersion = data['policyFormatVersion']
+        else:
+            raise CLIError(None, f"policyFormatVersion missing in Client Encryption Policy. Please verify your Client Encryption Policy JSON string")
+
+        if(policyFormatVersion < 1 or policyFormatVersion > 2):
+            raise CLIError(None, f"Invalid policyFormatVersion used in Client Encryption Policy. Please verify your Client Encryption Policy JSON string. Supported version are 1 and 2.")
+
+
+        listofIncludedPaths = []        
+        for includedpath in includedpaths:            
+            if "encryptionType" in includedpath:
+                encryptionType = includedpath['encryptionType']
+            else:
+                raise CLIError(None, f"encryptionType missing in includedPaths. Please verify your Client Encryption Policy JSON string")
+
+            if(encryptionType == ""):
+                raise CLIError(None, f"Invalid encryptionType included in Client Encryption Policy. encryptionType cannot be null or empty.")
+
+            if(encryptionType != "Deterministic" and encryptionType != "Randomized"):
+                raise CLIError(None, f"Invalid Encryption Type {encryptionType} used. Supported types are Deterministic or Randomized")
+
+            if "path" in includedpath:
+                path = includedpath['path']
+            else:
+                raise CLIError(None, f"path missing in includedPaths. Please verify your Client Encryption Policy JSON string")
+
+            if(path == ""):
+                raise CLIError(None, f"Invalid path included in Client Encryption Policy. Path cannot be null or empty.")
+
+            if(path[0] != "/" or path[-1] == "/"):
+                raise CLIError(None, 'Invalid path included in Client Encryption Policy. Only top level paths supported. Paths should begin with /. ')
+
+            if(path[1:] == "id"):
+                if(policyFormatVersion < 2):
+                    raise CLIError(None, f"id path which is part of Client Encryption policy is configured with invalid policyFormatVersion: {policyFormatVersion}. Please use policyFormatVersion 2.")
+
+                if(encryptionType != "Deterministic"):
+                    raise CLIError(None, f"id path is part of Client Encryption policy with invalid encryption type: {encryptionType}. Only deterministic encryption type is supported.")
+
+            if "clientEncryptionKeyId" in includedpath:
+                clientEncryptionKeyId = includedpath['clientEncryptionKeyId']
+            else:
+                raise CLIError(None, f"clientEncryptionKeyId missing in includedPaths. Please verify your Client Encryption Policy JSON string")
+
+            if(clientEncryptionKeyId == ""):
+                raise CLIError(None, f"Invalid clientEncryptionKeyId included in Client Encryption Policy. clientEncryptionKeyId cannot be null or empty.")
+
+            # for each partition key path verify if its part of client encryption policy or if its stop level path is part of client encryption policy
+            # eg: pk path is /a/b/c and /a is part of client encryption policy
+            for pkpath in sql_container_resource.partition_key.paths:
+                if(path[1:] == pkpath.split('/')[1]) and (encryptionType != "Deterministic"):
+                    if(policyFormatVersion < 2):
+                        raise CLIError(None, f"Partition key path:{pkpath} which is part of Client Encryption policy is configured with invalid policyFormatVersion: {policyFormatVersion}. Please use policyFormatVersion 2.")
+
+                    if(encryptionType != "Deterministic"):
+                        raise CLIError(None, 'Partition key path:{pkpath} is part of Client Encryption policy with invalid encryption type. Only deterministic encryption type is supported.')
+
+            if "encryptionAlgorithm" in includedpath:
+                encryptionAlgorithm = includedpath['encryptionAlgorithm']
+            else:
+                raise CLIError(None, f"encryptionAlgorithm missing in includedPaths. Please verify your Client Encryption Policy JSON string")
+
+            if(encryptionAlgorithm == ""):
+                raise CLIError(None, f"Invalid encryptionAlgorithm included in Client Encryption Policy. encryptionAlgorithm cannot be null or empty.")
+
+            if(encryptionAlgorithm != "AEAD_AES_256_CBC_HMAC_SHA256"):
+                raise CLIError(None, f"Invalid encryptionAlgorithm included in Client Encryption Policy. encryptionAlgorithm should be 'AEAD_AES_256_CBC_HMAC_SHA256'")
+
+            clientEncryptionIncludedPathObj = ClientEncryptionIncludedPath(path = path, client_encryption_key_id = clientEncryptionKeyId , encryption_type = encryptionType, encryption_algorithm = encryptionAlgorithm)
+            listofIncludedPaths.append(clientEncryptionIncludedPathObj)          
+
+        clientEncryptionPolicyObj = ClientEncryptionPolicy(included_paths = listofIncludedPaths, policy_format_version = policyFormatVersion)
+
+        # looks good set the client encryption policy object.
+        sql_container_resource.client_encryption_policy = clientEncryptionPolicyObj
 
 def cli_cosmosdb_sql_container_create(client,
                                       resource_group_name,
@@ -558,6 +653,7 @@ def cli_cosmosdb_sql_container_create(client,
                                       partition_key_version=None,
                                       default_ttl=None,
                                       indexing_policy=DEFAULT_INDEXING_POLICY,
+                                      client_encryption_policy=None,
                                       throughput=None,
                                       max_throughput=None,
                                       unique_key_policy=None,
@@ -571,10 +667,13 @@ def cli_cosmosdb_sql_container_create(client,
                                        default_ttl,
                                        indexing_policy,
                                        unique_key_policy,
+                                       client_encryption_policy,
                                        partition_key_version,
                                        conflict_resolution_policy,
                                        analytical_storage_ttl)
 
+    #print("Container ====== " + str(sql_container_resource))
+
     options = _get_options(throughput, max_throughput)
 
     sql_container_create_update_resource = SqlContainerCreateUpdateParameters(
@@ -606,6 +705,7 @@ def cli_cosmosdb_sql_container_update(client,
     sql_container_resource.default_ttl = sql_container.resource.default_ttl
     sql_container_resource.unique_key_policy = sql_container.resource.unique_key_policy
     sql_container_resource.conflict_resolution_policy = sql_container.resource.conflict_resolution_policy
+    sql_container_resource.client_encryption_policy = sql_container.resource.client_encryption_policy
 
     if _populate_sql_container_definition(sql_container_resource,
                                           None,
@@ -1959,7 +2059,8 @@ def cli_cosmosdb_collection_delete(client, database_id, collection_id):
 def _populate_collection_definition(collection,
                                     partition_key_path=None,
                                     default_ttl=None,
-                                    indexing_policy=None):
+                                    indexing_policy=None,
+                                    client_encryption_policy=None):
     if all(arg is None for arg in [partition_key_path, default_ttl, indexing_policy]):
         return False
 
@@ -1977,6 +2078,11 @@ def _populate_collection_definition(collection,
     if indexing_policy is not None:
         collection['indexingPolicy'] = indexing_policy
 
+    print("_populate_collection_definition: filling client encryption policy")
+
+    if client_encryption_policy is not None:
+        collection['clientEncryptionPolicy'] = client_encryption_policy
+
     return True
 
 
@@ -1986,7 +2092,8 @@ def cli_cosmosdb_collection_create(client,
                                    throughput=None,
                                    partition_key_path=None,
                                    default_ttl=None,
-                                   indexing_policy=DEFAULT_INDEXING_POLICY):
+                                   indexing_policy=DEFAULT_INDEXING_POLICY,
+                                   client_encryption_policy=None):
     """Creates an Azure Cosmos DB collection """
     collection = {'id': collection_id}
 
@@ -1997,8 +2104,10 @@ def cli_cosmosdb_collection_create(client,
     _populate_collection_definition(collection,
                                     partition_key_path,
                                     default_ttl,
-                                    indexing_policy)
+                                    indexing_policy,
+                                    client_encryption_policy)
 
+    #print("cli_cosmosdb_collection_create: filling client encryption policy ================  =======" + str(collection))
     created_collection = client.CreateContainer(_get_database_link(database_id), collection,
                                                 options)
     offer = _find_offer(client, created_collection['_self'])

src/azure-cli/requirements.py3.Darwin.txt

@@ -34,7 +34,7 @@ azure-mgmt-containerinstance==9.1.0
 azure-mgmt-containerregistry==8.2.0
 azure-mgmt-containerservice==19.1.0
 azure-mgmt-core==1.3.0
-azure-mgmt-cosmosdb==7.0.0b2
+azure-mgmt-cosmosdb==7.0.0b6
 azure-mgmt-databoxedge==1.0.0
 azure-mgmt-datalake-analytics==0.2.1
 azure-mgmt-datalake-nspkg==3.0.1

src/azure-cli/requirements.py3.Linux.txt

@@ -34,7 +34,7 @@ azure-mgmt-containerinstance==9.1.0
 azure-mgmt-containerregistry==8.2.0
 azure-mgmt-containerservice==19.1.0
 azure-mgmt-core==1.3.0
-azure-mgmt-cosmosdb==7.0.0b2
+azure-mgmt-cosmosdb==7.0.0b6
 azure-mgmt-databoxedge==1.0.0
 azure-mgmt-datalake-analytics==0.2.1
 azure-mgmt-datalake-nspkg==3.0.1

src/azure-cli/requirements.py3.windows.txt

@@ -34,7 +34,7 @@ azure-mgmt-containerinstance==9.1.0
 azure-mgmt-containerregistry==8.2.0
 azure-mgmt-containerservice==19.1.0
 azure-mgmt-core==1.3.0
-azure-mgmt-cosmosdb==7.0.0b2
+azure-mgmt-cosmosdb==7.0.0b6
 azure-mgmt-databoxedge==1.0.0
 azure-mgmt-datalake-analytics==0.2.1
 azure-mgmt-datalake-nspkg==3.0.1