|
25 | 25 | """ |
26 | 26 |
|
27 | 27 | import boto3 |
| 28 | +from aws_dbesdk_dynamodb.structures.dynamodb import LegacyPolicy |
28 | 29 |
|
29 | 30 | # Import from new AWS Database Encryption SDK |
30 | 31 | from .common import setup_awsdbe_client_with_legacy_override, setup_pure_awsdbe_client |
31 | 32 |
|
32 | 33 |
|
33 | | -def migration_step_2(kms_key_id, ddb_table_name): |
| 34 | +def migration_step_2(kms_key_id: str, ddb_table_name: str, sort_read_value: int=2): |
34 | 35 | """ |
35 | 36 | Migration Step 2: Using pure AWS DBESDK and legacy override together. |
36 | 37 |
|
37 | 38 | Args: |
38 | 39 | kms_key_id: The ARN of the KMS key to use for encryption |
39 | 40 | ddb_table_name: The name of the DynamoDB table |
| 41 | + sort_read_value: The sort key value to read |
40 | 42 |
|
41 | 43 | """ |
42 | | - ddb_client = boto3.client("dynamodb") |
43 | | - |
44 | | - # 1. Create two AWS DBESDK clients: |
45 | | - # - A pure client that can only read items encrypted with AWS DBESDK |
46 | | - # - A client with legacy override that can read legacy items |
47 | | - pure_client = setup_pure_awsdbe_client(kms_key_id, ddb_table_name) |
48 | | - legacy_override_client = setup_awsdbe_client_with_legacy_override(kms_key_id, ddb_table_name) |
49 | | - |
50 | | - # 2. Write a new item using the pure client (it will use the new encryption format) |
51 | | - item = { |
| 44 | + # 1. Create a DynamoDB Encryption SDK client with legacy override. |
| 45 | + # When configuring our legacy behavior, use `FORBID_LEGACY_ENCRYPT_ALLOW_LEGACY_DECRYPT`. |
| 46 | + # With this policy, you will continue to read items in both formats, |
| 47 | + # but will only write new items using the new format. |
| 48 | + encrypted_client = setup_awsdbe_client_with_legacy_override(kms_key_id, ddb_table_name, policy=LegacyPolicy.FORBID_LEGACY_ENCRYPT_ALLOW_LEGACY_DECRYPT) |
| 49 | + |
| 50 | + # 2. Put an item into your table using the DB ESDK Client. |
| 51 | + # This item will be encrypted in the latest format, using the |
| 52 | + # configuration from your modelled class to decide |
| 53 | + # which attribute to encrypt and/or sign. |
| 54 | + item_to_encrypt = { |
52 | 55 | "partition_key": {"S": "MigrationExampleForPython"}, |
53 | 56 | "sort_key": {"N": str(2)}, |
54 | 57 | "attribute1": {"S": "encrypt and sign me!"}, |
55 | 58 | "attribute2": {"S": "sign me!"}, |
56 | 59 | "attribute3": {"S": "ignore me!"}, |
57 | 60 | } |
58 | 61 |
|
59 | | - encrypted_item = pure_client.EncryptItem(item) |
60 | | - ddb_client.put_item(TableName=ddb_table_name, Item=encrypted_item) |
61 | | - print("Put item with sort_key=2 using pure AWS DBESDK client") |
62 | | - |
63 | | - # 3. Attempt to read a legacy-encrypted item (sort_key=0) with the pure client |
64 | | - # This should fail as the pure client doesn't understand the legacy format |
65 | | - key = {"partition_key": {"S": "MigrationExampleForPython"}, "sort_key": {"N": str(0)}} |
66 | | - |
67 | | - response = ddb_client.get_item(TableName=ddb_table_name, Key=key) |
68 | | - |
69 | | - if "Item" in response: |
70 | | - legacy_item = response["Item"] |
71 | | - try: |
72 | | - pure_client.DecryptItem(legacy_item) |
73 | | - print("Unexpectedly succeeded in decrypting legacy item with pure client!") |
74 | | - except Exception as e: |
75 | | - print(f"Expected failure: Pure client can't decrypt legacy item: {type(e).__name__}") |
76 | | - |
77 | | - # 4. Read the same legacy item using the client with legacy override |
78 | | - # This should succeed as this client understands both legacy and new formats |
79 | | - if "Item" in response: |
80 | | - legacy_item = response["Item"] |
81 | | - legacy_override_client.DecryptItem(legacy_item) |
82 | | - print("Successfully decrypted legacy item (sort_key=0) using client with legacy override") |
83 | | - |
84 | | - # 5. Read the new item (sort_key=2) with both clients |
85 | | - # Both should be able to read the new format |
86 | | - key = {"partition_key": {"S": "MigrationExampleForPython"}, "sort_key": {"N": str(2)}} |
87 | | - |
88 | | - response = ddb_client.get_item(TableName=ddb_table_name, Key=key) |
89 | | - |
90 | | - if "Item" in response: |
91 | | - new_item = response["Item"] |
92 | | - pure_client.DecryptItem(new_item) |
93 | | - print("Successfully decrypted new item (sort_key=2) using pure client") |
94 | | - |
95 | | - legacy_override_client.DecryptItem(new_item) |
96 | | - print("Successfully decrypted new item (sort_key=2) using client with legacy override") |
97 | | - |
98 | | - # 6. Re-encrypt a legacy item with the new format |
99 | | - # This is how you migrate your data from the legacy format to the new format |
100 | | - key = {"partition_key": {"S": "MigrationExampleForPython"}, "sort_key": {"N": str(0)}} |
101 | | - |
102 | | - response = ddb_client.get_item(TableName=ddb_table_name, Key=key) |
103 | | - |
104 | | - if "Item" in response: |
105 | | - # Get the legacy encrypted item |
106 | | - legacy_item = response["Item"] |
107 | | - |
108 | | - # Decrypt it with the client that has legacy override |
109 | | - decrypted_item = legacy_override_client.DecryptItem(legacy_item) |
110 | | - print("Successfully decrypted legacy item for re-encryption") |
111 | | - |
112 | | - # Re-encrypt with the pure client (new format) |
113 | | - new_encrypted_item = pure_client.EncryptItem(decrypted_item) |
114 | | - print("Successfully re-encrypted item in new format") |
115 | | - |
116 | | - # Store the re-encrypted item back in DynamoDB |
117 | | - ddb_client.put_item(TableName=ddb_table_name, Item=new_encrypted_item) |
118 | | - print("Successfully stored re-encrypted item") |
119 | | - |
120 | | - # Verify we can now read it with the pure client |
121 | | - verify_response = ddb_client.get_item(TableName=ddb_table_name, Key=key) |
122 | | - |
123 | | - if "Item" in verify_response: |
124 | | - re_encrypted_item = verify_response["Item"] |
125 | | - pure_client.DecryptItem(re_encrypted_item) |
126 | | - print("Successfully verified re-encrypted item can be read with pure client") |
127 | | - |
128 | | - |
129 | | -# def run_example(): |
130 | | -# """Run the full migration example.""" |
131 | | - |
132 | | -# print("\n=== Migration Example Step 2: Using Pure AWS DBESDK and Legacy Override Together ===") |
133 | | - |
134 | | -# try: |
135 | | -# # Ensure we have items from previous steps |
136 | | -# # First run migration step 0 to have legacy items |
137 | | -# migration_step0.migration_step0( |
138 | | -# kms_key_id=common.KMS_KEY_ID, |
139 | | -# ddb_table_name=common.TABLE_NAME |
140 | | -# ) |
141 | | - |
142 | | -# # Now run migration step 2 |
143 | | -# migration_step2( |
144 | | -# kms_key_id=common.KMS_KEY_ID, |
145 | | -# ddb_table_name=common.TABLE_NAME |
146 | | -# ) |
147 | | - |
148 | | -# print("\nMigration Step 2 completed successfully!") |
149 | | -# print("This demonstrates how to use both clients side-by-side and") |
150 | | -# print("how to re-encrypt legacy data to the new format.") |
| 62 | + put_item_request = { |
| 63 | + "TableName": ddb_table_name, |
| 64 | + "Item": item_to_encrypt, |
| 65 | + } |
151 | 66 |
|
152 | | -# except Exception as e: |
153 | | -# print(f"Error in Migration Step 2: {e}") |
154 | | -# raise |
| 67 | + put_item_response = encrypted_client.put_item(**put_item_request) |
| 68 | + # Demonstrate that PutItem succeeded |
| 69 | + assert put_item_response["ResponseMetadata"]["HTTPStatusCode"] == 200 |
| 70 | + print("Put Item Response SHOULD have legacy attributes") |
| 71 | + print(f"Put Item Response: {put_item_response['Item']}") |
| 72 | + |
| 73 | + # 3. Get an item back from the table using the Client. |
| 74 | + # If this is an item written in the old format (e.g. any item written |
| 75 | + # during Step 0 or 1), then we will attempt to decrypt the item |
| 76 | + # using the legacy behavior. |
| 77 | + # If this is an item written in the new format (e.g. any item written |
| 78 | + # during Step 2 or after), then we will attempt to decrypt the item using |
| 79 | + # the non-legacy behavior. |
| 80 | + key_to_get = {"partition_key": {"S": "MigrationExampleForPython"}, "sort_key": {"N": str(sort_read_value)}} |
| 81 | + |
| 82 | + get_item_request = {"TableName": ddb_table_name, "Key": key_to_get} |
| 83 | + get_item_response = encrypted_client.get_item(**get_item_request) |
| 84 | + |
| 85 | + # Demonstrate that GetItem succeeded and returned the decrypted item |
| 86 | + assert get_item_response["ResponseMetadata"]["HTTPStatusCode"] == 200 |
| 87 | + decrypted_item = get_item_response["Item"] |
| 88 | + # Demonstrate we get the expected item back |
| 89 | + assert decrypted_item["partition_key"]["S"] == "MigrationExampleForPython" |
| 90 | + assert decrypted_item["sort_key"]["N"] == str(sort_read_value) |
| 91 | + assert decrypted_item["attribute1"]["S"] == "encrypt and sign me!" |
| 92 | + assert decrypted_item["attribute2"]["S"] == "sign me!" |
| 93 | + assert decrypted_item["attribute3"]["S"] == "ignore me!" |
0 commit comments