[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

@@ -23,7 +23,8 @@
     validate_gossip_certificates,
     validate_client_certificates,
     validate_seednodes,
-    validate_node_count)
+    validate_node_count,
+    validate_client_encryption_policy)
 
 from azure.cli.command_modules.cosmosdb.actions import (
     CreateLocation, CreateDatabaseRestoreResource, UtcDatetimeAction, InvokeCommandArgumentsAddAction)
@@ -37,6 +38,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 +133,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(), validator=validate_client_encryption_policy, 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 +187,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(), validator=validate_client_encryption_policy, 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/_validators.py

@@ -248,3 +248,129 @@ def validate_node_count(ns):
     if ns.node_count is not None:
         if int(ns.node_count) < 3:
             raise InvalidArgumentValueError("""Node count cannot be less than 3.""")
+
+
+def validate_client_encryption_policy(ns):
+    """ Validate all the included paths in client encryption policy"""
+    partition_key_path = [ns.partition_key_path]
+    client_encryption_policy = ns.client_encryption_policy
+
+    if client_encryption_policy is not None:
+        if "includedPaths" in client_encryption_policy:
+            includedPaths = client_encryption_policy['includedPaths']
+        else:
+            raise InvalidArgumentValueError("includedPaths missing in Client Encryption Policy. "
+                                            "Please verify your Client Encryption Policy JSON string.")
+
+        if includedPaths == "":
+            raise InvalidArgumentValueError("includedPaths missing in Client Encryption Policy. "
+                                            "includedPaths cannot be null or empty.")
+
+        if "policyFormatVersion" in client_encryption_policy:
+            policyFormatVersion = client_encryption_policy['policyFormatVersion']
+        else:
+            raise InvalidArgumentValueError("policyFormatVersion missing in Client Encryption Policy. "
+                                            "Please verify your Client Encryption Policy JSON string.")
+
+        if not isinstance(policyFormatVersion, int):
+            raise InvalidArgumentValueError("Invalid policyFormatVersion type used in Client Encryption Policy. "
+                                            "policyFormatVersion is an integer type. "
+                                            "Supported versions are 1 and 2.")
+
+        if(policyFormatVersion < 1 or policyFormatVersion > 2):
+            raise InvalidArgumentValueError("Invalid policyFormatVersion used in Client Encryption Policy. "
+                                            "Please verify your Client Encryption Policy JSON string. "
+                                            "Supported versions are 1 and 2.")
+
+        _validate_included_paths_in_cep(partition_key_path, includedPaths, policyFormatVersion)
+
+
+def _validate_included_paths_in_cep(partition_key_path, includedPaths, policyFormatVersion):
+    listOfPaths = []
+    for includedPath in includedPaths:
+        if "encryptionType" in includedPath:
+            encryptionType = includedPath['encryptionType']
+        else:
+            raise InvalidArgumentValueError("encryptionType missing in includedPaths. "
+                                            "Please verify your Client Encryption Policy JSON string.")
+
+        if encryptionType == "":
+            raise InvalidArgumentValueError("Invalid encryptionType included in Client Encryption Policy. "
+                                            "encryptionType cannot be null or empty.")
+
+        if(encryptionType != "Deterministic" and encryptionType != "Randomized"):
+            raise InvalidArgumentValueError(f"Invalid Encryption Type {encryptionType} used. "
+                                            "Supported types are Deterministic or Randomized.")
+
+        if "path" in includedPath:
+            path = includedPath['path']
+        else:
+            raise InvalidArgumentValueError("path missing in includedPaths. "
+                                            "Please verify your Client Encryption Policy JSON string.")
+
+        if path == "":
+            raise InvalidArgumentValueError("Invalid path included in Client Encryption Policy. "
+                                            "Path cannot be null or empty.")
+
+        if path in listOfPaths:
+            raise InvalidArgumentValueError(f"Duplicate path:{path} found in Client Encryption Policy.")
+
+        listOfPaths.append(path)
+
+        if(path[0] != "/" or path.rfind('/') != 0):
+            raise InvalidArgumentValueError("Invalid path included in Client Encryption Policy. "
+                                            "Only top level paths supported. Paths should begin with /.")
+
+        if path[1:] == "id":
+            if policyFormatVersion < 2:
+                raise InvalidArgumentValueError("id path which is part of Client Encryption policy is configured "
+                                                f"with invalid policyFormatVersion: {policyFormatVersion}. "
+                                                "Please use policyFormatVersion 2.")
+
+            if encryptionType != "Deterministic":
+                raise InvalidArgumentValueError("id path is part of Client Encryption policy "
+                                                f"with invalid encryption type: {encryptionType}. "
+                                                "Only deterministic encryption type is supported.")
+
+        if "clientEncryptionKeyId" in includedPath:
+            clientEncryptionKeyId = includedPath['clientEncryptionKeyId']
+        else:
+            raise InvalidArgumentValueError("clientEncryptionKeyId missing in includedPaths. "
+                                            "Please verify your Client Encryption Policy JSON string.")
+
+        if clientEncryptionKeyId == "":
+            raise InvalidArgumentValueError("Invalid clientEncryptionKeyId included in Client Encryption Policy. "
+                                            "clientEncryptionKeyId cannot be null or empty.")
+
+        _validate_pk_paths_in_cep(path, partition_key_path, policyFormatVersion, encryptionType)
+
+        if "encryptionAlgorithm" in includedPath:
+            encryptionAlgorithm = includedPath['encryptionAlgorithm']
+        else:
+            raise InvalidArgumentValueError("encryptionAlgorithm missing in includedPaths. "
+                                            "Please verify your Client Encryption Policy JSON string.")
+
+        if encryptionAlgorithm == "":
+            raise InvalidArgumentValueError("Invalid encryptionAlgorithm included in Client Encryption Policy. "
+                                            "encryptionAlgorithm cannot be null or empty.")
+
+        if encryptionAlgorithm != "AEAD_AES_256_CBC_HMAC_SHA256":
+            raise InvalidArgumentValueError("Invalid encryptionAlgorithm included in Client Encryption Policy. "
+                                            "encryptionAlgorithm should be 'AEAD_AES_256_CBC_HMAC_SHA256'.")
+
+
+def _validate_pk_paths_in_cep(path, partition_key_path, policyFormatVersion, encryptionType):
+    """ for each partition key path verify if its part of client encryption policy
+    or if its top 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 partition_key_path:
+        if path[1:] == pkPath.split('/')[1]:
+            if policyFormatVersion < 2:
+                raise InvalidArgumentValueError(f"Partition key path:{pkPath} which is part of "
+                                                "Client Encryption policy is configured "
+                                                f"with invalid policyFormatVersion: {policyFormatVersion}. "
+                                                "Please use policyFormatVersion 2.")
+            if encryptionType != "Deterministic":
+                raise InvalidArgumentValueError(f"Partition key path:{pkPath} is part of "
+                                                "Client Encryption policy with invalid encryption type. "
+                                                "Only deterministic encryption type is supported.")

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

@@ -516,11 +516,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 +541,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:
+        sql_container_resource.client_encryption_policy = client_encryption_policy
+
     if conflict_resolution_policy is not None:
         sql_container_resource.conflict_resolution_policy = conflict_resolution_policy
 
@@ -558,6 +562,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 +576,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 +612,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 +621,7 @@ def cli_cosmosdb_sql_container_update(client,
                                           None,
                                           None,
                                           None,
+                                          None,
                                           analytical_storage_ttl):
         logger.debug('replacing SQL container')
 
@@ -1959,7 +1967,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 +1986,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 +1998,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 +2010,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)