[Cosmos DB] az cosmosdb sql container create: Add support to create containers with client encryption policy

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

@@ -37,6 +37,9 @@
 SQL_UNIQUE_KEY_POLICY_EXAMPLE = """--unique-key-policy "{\\"uniqueKeys\\": [{\\"paths\\": [\\"/path/to/key1\\"]}, {\\"paths\\": [\\"/path/to/key2\\"]}]}"
 """
 
+SQL_CLIENT_ENCRYPTION_POLICY_EXAMPLE = """--cep "{\\"includedPaths\\": [{\\"path\\": \\"/path1\\",\\"clientEncryptionKeyId\\": \\"key1\\",\\"encryptionAlgorithm\\": \\"AEAD_AES_256_CBC_HMAC_SHA256\\",\\"encryptionType\\": \\"Deterministic\\"}],\\"policyFormatVersion\\": 2}"
+"""
+
 SQL_GREMLIN_CONFLICT_RESOLUTION_POLICY_EXAMPLE = """--conflict-resolution-policy "{\\"mode\\": \\"lastWriterWins\\", \\"conflictResolutionPath\\": \\"/path\\"}"
 """
 
@@ -129,6 +132,7 @@ def load_arguments(self, _):
         c.argument('collection_id', options_list=['--collection-name', '-c'], help='Collection Name')
         c.argument('throughput', type=int, help='Offer Throughput (RU/s)')
         c.argument('partition_key_path', help='Partition Key Path, e.g., \'/properties/name\'')
+        c.argument('client_encryption_policy', options_list=['--cep'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Client Encryption Policy, you can enter it as a string or as a file, e.g., --cep @policy-file.json or ' + SQL_CLIENT_ENCRYPTION_POLICY_EXAMPLE)
         c.argument('indexing_policy', type=shell_safe_json_parse, completer=FilesCompleter(), help='Indexing Policy, you can enter it as a string or as a file, e.g., --indexing-policy @policy-file.json)')
         c.argument('default_ttl', type=int, help='Default TTL. Provide 0 to disable.')
 
@@ -182,6 +186,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=['--cep'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Client Encryption Policy, you can enter it as a string or as a file, e.g., --cep @policy-file.json or ' + SQL_CLIENT_ENCRYPTION_POLICY_EXAMPLE)
         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:
@@ -540,6 +544,9 @@ def _populate_sql_container_definition(sql_container_resource,
     if unique_key_policy is not None:
         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
 
@@ -549,6 +556,109 @@ 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 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, "includedPaths missing in Client Encryption Policy. Please verify your Client Encryption Policy JSON string")
+
+        if includedPaths == "":
+            raise CLIError(None, "includedPaths missing in Client Encryption Policy. includedPaths cannot be null or empty.")
+
+        if "policyFormatVersion" in data:
+            policyFormatVersion = data['policyFormatVersion']
+        else:
+            raise CLIError(None, "policyFormatVersion missing in Client Encryption Policy. Please verify your Client Encryption Policy JSON string")
+
+        if not isinstance(policyFormatVersion, int):
+            raise CLIError(None, "Invalid policyFormatVersion format in Client Encryption Policy. policyFormatVersion is an integer type. Supported version are 1 and 2.")
+
+        if(policyFormatVersion < 1 or policyFormatVersion > 2):
+            raise CLIError(None, "Invalid policyFormatVersion used in Client Encryption Policy. Please verify your Client Encryption Policy JSON string. Supported version are 1 and 2.")
+
+        listOfIncludedPaths = []
+        listOfPaths = []
+        for includedPath in includedPaths:
+            if "encryptionType" in includedPath:
+                encryptionType = includedPath['encryptionType']
+            else:
+                raise CLIError(None, "encryptionType missing in includedPaths. Please verify your Client Encryption Policy JSON string")
+
+            if encryptionType == "":
+                raise CLIError(None, "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, "path missing in includedPaths. Please verify your Client Encryption Policy JSON string")
+
+            if path in listOfPaths:
+                raise CLIError(None, f"Duplicate path:{path} found in Client Encryption Policy")
+
+            listOfPaths.append(path)
+
+            if path == "":
+                raise CLIError(None, "Invalid path included in Client Encryption Policy. Path cannot be null or empty.")
+
+            if(path[0] != "/" or path.rfind('/') != 0):
+                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, "clientEncryptionKeyId missing in includedPaths. Please verify your Client Encryption Policy JSON string")
+
+            if clientEncryptionKeyId == "":
+                raise CLIError(None, "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, f"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, "encryptionAlgorithm missing in includedPaths. Please verify your Client Encryption Policy JSON string")
+
+            if encryptionAlgorithm == "":
+                raise CLIError(None, "Invalid encryptionAlgorithm included in Client Encryption Policy. encryptionAlgorithm cannot be null or empty.")
+
+            if encryptionAlgorithm != "AEAD_AES_256_CBC_HMAC_SHA256":
+                raise CLIError(None, "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,
                                       account_name,
@@ -558,6 +668,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,6 +682,7 @@ 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)
@@ -606,6 +718,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,
@@ -614,6 +727,7 @@ def cli_cosmosdb_sql_container_update(client,
                                           None,
                                           None,
                                           None,
+                                          None,
                                           analytical_storage_ttl):
         logger.debug('replacing SQL container')
 
@@ -1959,7 +2073,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 +2092,9 @@ def _populate_collection_definition(collection,
     if indexing_policy is not None:
         collection['indexingPolicy'] = indexing_policy
 
+    if client_encryption_policy is not None:
+        collection['clientEncryptionPolicy'] = client_encryption_policy
+
     return True
 
 
@@ -1986,7 +2104,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,7 +2116,8 @@ def cli_cosmosdb_collection_create(client,
     _populate_collection_definition(collection,
                                     partition_key_path,
                                     default_ttl,
-                                    indexing_policy)
+                                    indexing_policy,
+                                    client_encryption_policy)
 
     created_collection = client.CreateContainer(_get_database_link(database_id), collection,
                                                 options)