chore(python): examples for searchable encryption

Author: imabhichow

Examples/runtimes/python/DynamoDBEncryption/src/get_encrypted_data_key_description_example.py

@@ -0,0 +1,82 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Example demonstrating how to get encrypted data key descriptions from DynamoDB items."""
+
+import boto3
+from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb.client import DynamoDbEncryption
+from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb.config import (
+    DynamoDbEncryptionConfig,
+)
+from aws_dbesdk_dynamodb.structures.dynamodb import (
+    GetEncryptedDataKeyDescriptionInput,
+    GetEncryptedDataKeyDescriptionUnionItem,
+)
+
+
+def get_encrypted_data_key_description(
+    table_name: str,
+    partition_key: str,
+    partition_key_val: str,
+    sort_key_name: str,
+    sort_key_value: str,
+    expected_key_provider_id: str,
+    expected_key_provider_info: str,
+    expected_branch_key_id: str,
+    expected_branch_key_version: str,
+):
+    """
+    Get encrypted data key description from a DynamoDB item.
+
+    :param table_name: The name of the DynamoDB table
+    :param partition_key: The name of the partition key
+    :param partition_key_val: The value of the partition key
+    :param sort_key_name: The name of the sort key
+    :param sort_key_value: The value of the sort key
+    :param expected_key_provider_id: The expected key provider ID
+    :param expected_key_provider_info: The expected key provider info (optional)
+    :param expected_branch_key_id: The expected branch key ID (optional)
+    :param expected_branch_key_version: The expected branch key version (optional)
+    """
+    # 1. Create a new AWS SDK DynamoDb client. This client will be used to get item from the DynamoDB table
+    ddb = boto3.client("dynamodb")
+
+    # 2. Get item from the DynamoDB table. This item will be used to Get Encrypted DataKey Description
+    key_to_get = {partition_key: {"S": partition_key_val}, sort_key_name: {"N": sort_key_value}}
+
+    response = ddb.get_item(TableName=table_name, Key=key_to_get)
+
+    returned_item = response.get("Item", {})
+    if not returned_item:
+        print(f"No item found with the key {partition_key}!")
+        return
+
+    # 3. Prepare the input for GetEncryptedDataKeyDescription method.
+    # This input can be a DynamoDB item or a header. For now, we are giving input as a DynamoDB item
+    # but users can also extract the header from the attribute "aws_dbe_head" in the DynamoDB table
+    # and use it for GetEncryptedDataKeyDescription method.
+    ddb_enc = DynamoDbEncryption(config=DynamoDbEncryptionConfig())
+
+    input_union = GetEncryptedDataKeyDescriptionUnionItem(returned_item)
+
+    input_obj = GetEncryptedDataKeyDescriptionInput(input=input_union)
+
+    output = ddb_enc.get_encrypted_data_key_description(input=input_obj)
+
+    # In the following code, we are giving input as header instead of a complete DynamoDB item
+    # This code is provided solely to demo how the alternative approach works. So, it is commented.
+
+    # header_attribute = "aws_dbe_head"
+    # header = returned_item[header_attribute]["B"]
+    # input_union = GetEncryptedDataKeyDescriptionUnion(
+    #     header=header
+    # )
+
+    # Assert everything
+    assert output.encrypted_data_key_description_output[0].key_provider_id == expected_key_provider_id
+
+    if expected_key_provider_id.startswith("aws-kms"):
+        assert output.encrypted_data_key_description_output[0].key_provider_info == expected_key_provider_info
+
+    if output.encrypted_data_key_description_output[0].key_provider_id == "aws-kms-hierarchy":
+        assert output.encrypted_data_key_description_output[0].branch_key_id == expected_branch_key_id
+        assert output.encrypted_data_key_description_output[0].branch_key_version == expected_branch_key_version

Examples/runtimes/python/DynamoDBEncryption/src/searchable_encryption/basic_searchable_encryption_example.py

@@ -0,0 +1,298 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Example demonstrating DynamoDB encryption using beacon styles.
+
+This example demonstrates how to use Beacons Styles on Standard Beacons on encrypted attributes,
+    put an item with the beacon, and query against that beacon.
+This example follows a use case of a database that stores food information.
+    This is an extension of the "BasicSearchableEncryptionExample" in this directory
+    and uses the same table schema.
+
+Running this example requires access to a DDB table with the
+following key configuration:
+  - Partition key is named "work_id" with type (S)
+  - Sort key is named "inspection_time" with type (S)
+
+In this example for storing food information, this schema is utilized for the data:
+ - "work_id" stores a unique identifier for a unit inspection work order (v4 UUID)
+ - "inspection_date" stores an ISO 8601 date for the inspection (YYYY-MM-DD)
+ - "fruit" stores one type of fruit
+ - "basket" stores a set of types of fruit
+ - "dessert" stores one type of dessert
+ - "veggies" stores a set of types of vegetable
+ - "work_type" stores a unit inspection category
+
+The example requires the following ordered input command line parameters:
+  1. DDB table name for table to put/query data from
+  2. Branch key ID for a branch key that was previously created in your key store. See the
+     CreateKeyStoreKeyExample.
+  3. Branch key wrapping KMS key ARN for the KMS key used to create the branch key with ID
+     provided in arg 2
+  4. Branch key DDB table name for the DDB table representing the branch key store
+"""
+from typing import List
+
+import boto3
+from aws_cryptographic_material_providers.keystore.client import KeyStore
+from aws_cryptographic_material_providers.keystore.config import KeyStoreConfig
+from aws_cryptographic_material_providers.keystore.models import KMSConfigurationKmsKeyArn
+from aws_cryptographic_material_providers.mpl import AwsCryptographicMaterialProviders
+from aws_cryptographic_material_providers.mpl.config import MaterialProvidersConfig
+from aws_cryptographic_material_providers.mpl.models import CreateAwsKmsHierarchicalKeyringInput
+from aws_dbesdk_dynamodb.encrypted.client import EncryptedClient
+from aws_dbesdk_dynamodb.structures.dynamodb import (
+    AsSet,
+    BeaconKeySourceSingle,
+    BeaconStyleAsSet,
+    BeaconStylePartOnly,
+    BeaconStyleShared,
+    BeaconStyleSharedSet,
+    BeaconVersion,
+    CompoundBeacon,
+    DynamoDbTableEncryptionConfig,
+    DynamoDbTablesEncryptionConfig,
+    EncryptedPart,
+    PartOnly,
+    SearchConfig,
+    Shared,
+    SharedSet,
+    SignedPart,
+    SingleKeyStore,
+    StandardBeacon,
+)
+from aws_dbesdk_dynamodb.structures.structured_encryption import CryptoAction
+
+
+def put_item_query_item_with_beacon_styles(
+    ddb_table_name: str, branch_key_id: str, branch_key_wrapping_kms_key_arn: str, branch_key_ddb_table_name: str
+):
+    """
+    Demonstrate using beacon styles with DynamoDB encryption.
+
+    :param ddb_table_name: The name of the DynamoDB table
+    :param branch_key_id: Branch key ID for a branch key previously created in key store
+    :param branch_key_wrapping_kms_key_arn: ARN of KMS key used to create the branch key
+    :param branch_key_ddb_table_name: Name of DDB table representing the branch key store
+    """
+    # 1. Create Beacons.
+    standard_beacon_list: List[StandardBeacon] = []
+
+    # The fruit beacon allows searching on the encrypted fruit attribute
+    # We have selected 30 as an example beacon length, but you should go to
+    # https://docs.aws.amazon.com/database-encryption-sdk/latest/devguide/choosing-beacon-length.html
+    # when creating your beacons.
+    fruit_beacon = StandardBeacon(name="fruit", length=30)
+    standard_beacon_list.append(fruit_beacon)
+
+    # The basket beacon allows searching on the encrypted basket attribute
+    # basket is used as a Set, and therefore needs a beacon style to reflect that.
+    # Further, we need to be able to compare the items in basket to the fruit attribute
+    # so we `share` this beacon with `fruit`.
+    # Since we need both of these things, we use the SharedSet style.
+    basket_beacon = StandardBeacon(name="basket", length=30, style=BeaconStyleSharedSet(SharedSet(other="fruit")))
+    standard_beacon_list.append(basket_beacon)
+
+    # The dessert beacon allows searching on the encrypted dessert attribute
+    # We need to be able to compare the dessert attribute to the fruit attribute
+    # so we `share` this beacon with `fruit`.
+    dessert_beacon = StandardBeacon(name="dessert", length=30, style=BeaconStyleShared(Shared(other="fruit")))
+    standard_beacon_list.append(dessert_beacon)
+
+    # The veggie_beacon allows searching on the encrypted veggies attribute
+    # veggies is used as a Set, and therefore needs a beacon style to reflect that.
+    veggie_beacon = StandardBeacon(name="veggies", length=30, style=BeaconStyleAsSet(AsSet()))
+    standard_beacon_list.append(veggie_beacon)
+
+    # The work_type_beacon allows searching on the encrypted work_type attribute
+    # We only use it as part of the compound work_unit beacon,
+    # so we disable its use as a standalone beacon
+    work_type_beacon = StandardBeacon(name="work_type", length=30, style=BeaconStylePartOnly(PartOnly()))
+    standard_beacon_list.append(work_type_beacon)
+
+    # Here we build a compound beacon from work_id and work_type
+    # If we had tried to make a StandardBeacon from work_type, we would have seen an error
+    # because work_type is "PartOnly"
+    encrypted_part_list = [EncryptedPart(name="work_type", prefix="T-")]
+
+    signed_part_list = [SignedPart(name="work_id", prefix="I-")]
+
+    compound_beacon_list = [
+        CompoundBeacon(name="work_unit", split=".", encrypted=encrypted_part_list, signed=signed_part_list)
+    ]
+
+    # 2. Configure the Keystore
+    #    These are the same constructions as in the Basic example, which describes these in more detail.
+    keystore = KeyStore(
+        config=KeyStoreConfig(
+            ddb_client=boto3.client("dynamodb"),
+            ddb_table_name=branch_key_ddb_table_name,
+            logical_key_store_name=branch_key_ddb_table_name,
+            kms_client=boto3.client("kms"),
+            kms_configuration=KMSConfigurationKmsKeyArn(value=branch_key_wrapping_kms_key_arn),
+        )
+    )
+
+    # 3. Create BeaconVersion.
+    #    This is similar to the Basic example
+    beacon_versions = [
+        BeaconVersion(
+            standard_beacons=standard_beacon_list,
+            compound_beacons=compound_beacon_list,
+            version=1,  # MUST be 1
+            key_store=keystore,
+            key_source=BeaconKeySourceSingle(SingleKeyStore(key_id=branch_key_id, cache_ttl=6000)),
+        )
+    ]
+
+    # 4. Create a Hierarchical Keyring
+    #    This is the same configuration as in the Basic example.
+    mat_prov = AwsCryptographicMaterialProviders(config=MaterialProvidersConfig())
+
+    keyring_input = CreateAwsKmsHierarchicalKeyringInput(
+        branch_key_id=branch_key_id, key_store=keystore, ttl_seconds=6000
+    )
+
+    kms_keyring = mat_prov.create_aws_kms_hierarchical_keyring(input=keyring_input)
+
+    # 5. Configure which attributes are encrypted and/or signed when writing new items.
+    attribute_actions = {
+        "work_id": CryptoAction.SIGN_ONLY,  # Our partition attribute must be SIGN_ONLY
+        "inspection_date": CryptoAction.SIGN_ONLY,  # Our sort attribute must be SIGN_ONLY
+        "dessert": CryptoAction.ENCRYPT_AND_SIGN,  # Beaconized attributes must be encrypted
+        "fruit": CryptoAction.ENCRYPT_AND_SIGN,  # Beaconized attributes must be encrypted
+        "basket": CryptoAction.ENCRYPT_AND_SIGN,  # Beaconized attributes must be encrypted
+        "veggies": CryptoAction.ENCRYPT_AND_SIGN,  # Beaconized attributes must be encrypted
+        "work_type": CryptoAction.ENCRYPT_AND_SIGN,  # Beaconized attributes must be encrypted
+    }
+
+    # 6. Create the DynamoDb Encryption configuration for the table we will be writing to.
+    #    The beaconVersions are added to the search configuration.
+    table_config = DynamoDbTableEncryptionConfig(
+        logical_table_name=ddb_table_name,
+        partition_key_name="work_id",
+        sort_key_name="inspection_date",
+        attribute_actions_on_encrypt=attribute_actions,
+        keyring=kms_keyring,
+        search=SearchConfig(write_version=1, versions=beacon_versions),  # MUST be 1
+    )
+
+    table_configs = {ddb_table_name: table_config}
+    tables_config = DynamoDbTablesEncryptionConfig(table_encryption_configs=table_configs)
+
+    # 7. Create the EncryptedClient
+    ddb_client = boto3.client("dynamodb")
+    encrypted_ddb_client = EncryptedClient(client=ddb_client, encryption_config=tables_config)
+
+    # 8. Create item one, specifically with "dessert != fruit", and "fruit in basket".
+    item1 = {
+        "work_id": {"S": "1"},
+        "inspection_date": {"S": "2023-06-13"},
+        "dessert": {"S": "cake"},
+        "fruit": {"S": "banana"},
+        "basket": {"SS": ["apple", "banana", "pear"]},
+        "veggies": {"SS": ["beans", "carrots", "celery"]},
+        "work_type": {"S": "small"},
+    }
+
+    # 9. Create item two, specifically with "dessert == fruit", and "fruit not in basket".
+    item2 = {
+        "work_id": {"S": "2"},
+        "inspection_date": {"S": "2023-06-13"},
+        "fruit": {"S": "orange"},
+        "dessert": {"S": "orange"},
+        "basket": {"SS": ["blackberry", "blueberry", "strawberry"]},
+        "veggies": {"SS": ["beans", "carrots", "peas"]},
+        "work_type": {"S": "large"},
+    }
+
+    # 10. Add the two items
+    put_response = encrypted_ddb_client.put_item(TableName=ddb_table_name, Item=item1)
+    # Validate object put successfully
+    assert put_response["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+    put_response = encrypted_ddb_client.put_item(TableName=ddb_table_name, Item=item2)
+    # Validate object put successfully
+    assert put_response["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+    # 11. Test the first type of Set operation:
+    # Select records where the basket attribute holds a particular value
+    expression_attribute_values = {":value": {"S": "banana"}}
+
+    scan_response = encrypted_ddb_client.scan(
+        TableName=ddb_table_name,
+        FilterExpression="contains(basket, :value)",
+        ExpressionAttributeValues=expression_attribute_values,
+    )
+    # Validate query was returned successfully
+    assert scan_response["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+    # Validate only 1 item was returned: item1
+    assert len(scan_response["Items"]) == 1
+    assert scan_response["Items"][0] == item1
+
+    # 12. Test the second type of Set operation:
+    # Select records where the basket attribute holds the fruit attribute
+    scan_response = encrypted_ddb_client.scan(TableName=ddb_table_name, FilterExpression="contains(basket, fruit)")
+    # Validate query was returned successfully
+    assert scan_response["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+    # Validate only 1 item was returned: item1
+    assert len(scan_response["Items"]) == 1
+    assert scan_response["Items"][0] == item1
+
+    # 13. Test the third type of Set operation:
+    # Select records where the fruit attribute exists in a particular set
+    expression_attribute_values = {":value": {"SS": ["boysenberry", "grape", "orange"]}}
+
+    scan_response = encrypted_ddb_client.scan(
+        TableName=ddb_table_name,
+        FilterExpression="contains(:value, fruit)",
+        ExpressionAttributeValues=expression_attribute_values,
+    )
+    # Validate query was returned successfully
+    assert scan_response["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+    # Validate only 1 item was returned: item2
+    assert len(scan_response["Items"]) == 1
+    assert scan_response["Items"][0] == item2
+
+    # 14. Test a Shared search. Select records where the dessert attribute matches the fruit attribute
+    scan_response = encrypted_ddb_client.scan(TableName=ddb_table_name, FilterExpression="dessert = fruit")
+    # Validate query was returned successfully
+    assert scan_response["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+    # Validate only 1 item was returned: item2
+    assert len(scan_response["Items"]) == 1
+    assert scan_response["Items"][0] == item2
+
+    # 15. Test the AsSet attribute 'veggies':
+    # Select records where the veggies attribute holds a particular value
+    expression_attribute_values = {":value": {"S": "peas"}}
+
+    scan_response = encrypted_ddb_client.scan(
+        TableName=ddb_table_name,
+        FilterExpression="contains(veggies, :value)",
+        ExpressionAttributeValues=expression_attribute_values,
+    )
+    # Validate query was returned successfully
+    assert scan_response["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+    # Validate only 1 item was returned: item2
+    assert len(scan_response["Items"]) == 1
+    assert scan_response["Items"][0] == item2
+
+    # 16. Test the compound beacon 'work_unit':
+    expression_attribute_values = {":value": {"S": "I-1.T-small"}}
+
+    scan_response = encrypted_ddb_client.scan(
+        TableName=ddb_table_name,
+        FilterExpression="work_unit = :value",
+        ExpressionAttributeValues=expression_attribute_values,
+    )
+    # Validate query was returned successfully
+    assert scan_response["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+    # Validate only 1 item was returned: item1
+    assert len(scan_response["Items"]) == 1
+    assert scan_response["Items"][0] == item1