Fix EncryptedPaginator to successfully decrypt

CHANGELOG.rst

@@ -2,6 +2,14 @@
 Changelog
 *********
 
+1.1.1 -- 2019-08-xx
+===================
+
+Bugfixes
+--------
+* Fix :class:`EncryptedPaginator` to successfully decrypt when using :class:`AwsKmsCryptographicMaterialsProvider`
+  `#118 <https://github.com/aws/aws-dynamodb-encryption-python/pull/118>`_
+
 1.1.0 -- 2019-03-13
 ===================
 

src/dynamodb_encryption_sdk/encrypted/client.py

@@ -22,6 +22,7 @@
     crypto_config_from_kwargs,
     decrypt_batch_get_item,
     decrypt_get_item,
+    decrypt_list_of_items,
     decrypt_multi_get,
     encrypt_batch_write_item,
     encrypt_put_item,
@@ -104,8 +105,11 @@ def paginate(self, **kwargs):
         crypto_config, ddb_kwargs = self._crypto_config_method(**kwargs)
 
         for page in self._paginator.paginate(**ddb_kwargs):
-            for pos, value in enumerate(page["Items"]):
-                page["Items"][pos] = self._decrypt_method(item=value, crypto_config=crypto_config)
+            page["Items"] = list(
+                decrypt_list_of_items(
+                    crypto_config=crypto_config, decrypt_method=self._decrypt_method, items=page["Items"]
+                )
+            )
             yield page
 
 

src/dynamodb_encryption_sdk/internal/utils.py

@@ -29,7 +29,7 @@
 from dynamodb_encryption_sdk.transform import dict_to_ddb
 
 try:  # Python 3.5.0 and 3.5.1 have incompatible typing modules
-    from typing import Any, Bool, Callable, Dict, Text  # noqa pylint: disable=unused-import
+    from typing import Any, Bool, Callable, Dict, Iterable, Text  # noqa pylint: disable=unused-import
 except ImportError:  # pragma: no cover
     # We only actually need these imports when running the mypy checks
     pass
@@ -41,6 +41,7 @@
     "crypto_config_from_cache",
     "decrypt_get_item",
     "decrypt_multi_get",
+    "decrypt_list_of_items",
     "decrypt_batch_get_item",
     "encrypt_put_item",
     "encrypt_batch_write_item",
@@ -171,6 +172,22 @@ def _item_transformer(crypto_transformer):
     return lambda x: x
 
 
+def decrypt_list_of_items(crypto_config, decrypt_method, items):
+    # type: (CryptoConfig, Callable, Iterable[Any]) -> Iterable[Any]
+    # TODO: narrow this down
+    """Iterate through a list of encrypted items, decrypting each item and yielding the plaintext item.
+
+    :param CryptoConfig crypto_config: :class:`CryptoConfig` to use
+    :param callable decrypt_method: Method to use to decrypt items
+    :param items: Iterable of encrypted items
+    :return: Iterable of plaintext items
+    """
+    for value in items:
+        yield decrypt_method(
+            item=value, crypto_config=crypto_config.with_item(_item_transformer(decrypt_method)(value))
+        )
+
+
 def decrypt_multi_get(decrypt_method, crypto_config_method, read_method, **kwargs):
     # type: (Callable, Callable, Callable, **Any) -> Dict
     # TODO: narrow this down
@@ -186,11 +203,9 @@ def decrypt_multi_get(decrypt_method, crypto_config_method, read_method, **kwarg
     validate_get_arguments(kwargs)
     crypto_config, ddb_kwargs = crypto_config_method(**kwargs)
     response = read_method(**ddb_kwargs)
-    for pos in range(len(response["Items"])):
-        response["Items"][pos] = decrypt_method(
-            item=response["Items"][pos],
-            crypto_config=crypto_config.with_item(_item_transformer(decrypt_method)(response["Items"][pos])),
-        )
+    response["Items"] = list(
+        decrypt_list_of_items(crypto_config=crypto_config, decrypt_method=decrypt_method, items=response["Items"])
+    )
     return response
 
 

test/functional/encrypted/test_client.py

@@ -20,7 +20,7 @@
     build_static_jce_cmp,
     client_batch_items_unprocessed_check,
     client_cycle_batch_items_check,
-    client_cycle_batch_items_check_paginators,
+    client_cycle_batch_items_check_scan_paginator,
     client_cycle_single_item_check,
     set_parametrized_actions,
     set_parametrized_cmp,
@@ -49,8 +49,8 @@ def _client_cycle_batch_items_check(materials_provider, initial_actions, initial
     )
 
 
-def _client_cycle_batch_items_check_paginators(materials_provider, initial_actions, initial_item):
-    return client_cycle_batch_items_check_paginators(
+def _client_cycle_batch_items_check_scan_paginator(materials_provider, initial_actions, initial_item):
+    return client_cycle_batch_items_check_scan_paginator(
         materials_provider, initial_actions, initial_item, TEST_TABLE_NAME, "us-west-2"
     )
 
@@ -71,9 +71,9 @@ def test_ephemeral_batch_item_cycle(example_table, some_cmps, parametrized_actio
     _client_cycle_batch_items_check(some_cmps, parametrized_actions, parametrized_item)
 
 
-def test_ephemeral_batch_item_cycle_paginators(example_table, some_cmps, parametrized_actions, parametrized_item):
-    """Test a small number of curated CMPs against a small number of curated items using paginators."""
-    _client_cycle_batch_items_check_paginators(some_cmps, parametrized_actions, parametrized_item)
+def test_ephemeral_batch_item_cycle_scan_paginator(example_table, some_cmps, parametrized_actions, parametrized_item):
+    """Test a small number of curated CMPs against a small number of curated items using the scan paginator."""
+    _client_cycle_batch_items_check_scan_paginator(some_cmps, parametrized_actions, parametrized_item)
 
 
 def test_batch_item_unprocessed(example_table, parametrized_actions, parametrized_item):

test/functional/functional_test_utils.py

@@ -437,30 +437,34 @@ def cycle_batch_item_check(
     check_attribute_actions = initial_actions.copy()
     check_attribute_actions.set_index_keys(*list(TEST_KEY.keys()))
     items = _generate_items(initial_item, write_transformer)
+    items_in_table = len(items)
 
     _put_result = encrypted.batch_write_item(  # noqa
         RequestItems={table_name: [{"PutRequest": {"Item": _item}} for _item in items]}
     )
 
-    ddb_keys = [write_transformer(key) for key in TEST_BATCH_KEYS]
-    encrypted_result = raw.batch_get_item(RequestItems={table_name: {"Keys": ddb_keys}})
-    check_many_encrypted_items(
-        actual=encrypted_result["Responses"][table_name],
-        expected=items,
-        attribute_actions=check_attribute_actions,
-        transformer=read_transformer,
-    )
-
-    decrypted_result = encrypted.batch_get_item(RequestItems={table_name: {"Keys": ddb_keys}})
-    assert_equal_lists_of_items(
-        actual=decrypted_result["Responses"][table_name], expected=items, transformer=read_transformer
-    )
+    try:
+        ddb_keys = [write_transformer(key) for key in TEST_BATCH_KEYS]
+        encrypted_result = raw.batch_get_item(RequestItems={table_name: {"Keys": ddb_keys}})
+        check_many_encrypted_items(
+            actual=encrypted_result["Responses"][table_name],
+            expected=items,
+            attribute_actions=check_attribute_actions,
+            transformer=read_transformer,
+        )
 
-    if delete_items:
-        _cleanup_items(encrypted, write_transformer, table_name)
+        decrypted_result = encrypted.batch_get_item(RequestItems={table_name: {"Keys": ddb_keys}})
+        assert_equal_lists_of_items(
+            actual=decrypted_result["Responses"][table_name], expected=items, transformer=read_transformer
+        )
+    finally:
+        if delete_items:
+            _cleanup_items(encrypted, write_transformer, table_name)
+            items_in_table = 0
 
     del check_attribute_actions
     del items
+    return items_in_table
 
 
 def cycle_batch_writer_check(raw_table, encrypted_table, initial_actions, initial_item):
@@ -692,16 +696,23 @@ def client_batch_items_unprocessed_check(
         )
 
 
-def client_cycle_batch_items_check_paginators(
+def client_cycle_batch_items_check_scan_paginator(
     materials_provider, initial_actions, initial_item, table_name, region_name=None
 ):
+    """Helper function for testing the "scan" paginator.
+
+    Populate the specified table with encrypted items,
+    scan the table with raw client paginator to get encrypted items,
+    scan the table with encrypted client paginator to get decrypted items,
+    then verify that all items appear to have been encrypted correctly.
+    """
     kwargs = {}
     if region_name is not None:
         kwargs["region_name"] = region_name
     client = boto3.client("dynamodb", **kwargs)
     e_client = EncryptedClient(client=client, materials_provider=materials_provider, attribute_actions=initial_actions)
 
-    cycle_batch_item_check(
+    items_in_table = cycle_batch_item_check(
         raw=client,
         encrypted=e_client,
         initial_actions=initial_actions,
@@ -712,29 +723,31 @@ def client_cycle_batch_items_check_paginators(
         delete_items=False,
     )
 
-    encrypted_items = []
-    raw_paginator = client.get_paginator("scan")
-    for page in raw_paginator.paginate(TableName=table_name, ConsistentRead=True):
-        encrypted_items.extend(page["Items"])
-
-    decrypted_items = []
-    encrypted_paginator = e_client.get_paginator("scan")
-    for page in encrypted_paginator.paginate(TableName=table_name, ConsistentRead=True):
-        decrypted_items.extend(page["Items"])
-
-    print(encrypted_items)
-    print(decrypted_items)
-
-    check_attribute_actions = initial_actions.copy()
-    check_attribute_actions.set_index_keys(*list(TEST_KEY.keys()))
-    check_many_encrypted_items(
-        actual=encrypted_items,
-        expected=decrypted_items,
-        attribute_actions=check_attribute_actions,
-        transformer=ddb_to_dict,
-    )
+    try:
+        encrypted_items = []
+        raw_paginator = client.get_paginator("scan")
+        for page in raw_paginator.paginate(TableName=table_name, ConsistentRead=True):
+            encrypted_items.extend(page["Items"])
+
+        decrypted_items = []
+        encrypted_paginator = e_client.get_paginator("scan")
+        for page in encrypted_paginator.paginate(TableName=table_name, ConsistentRead=True):
+            decrypted_items.extend(page["Items"])
+
+        assert encrypted_items and decrypted_items
+        assert len(encrypted_items) == len(decrypted_items) == items_in_table
+
+        check_attribute_actions = initial_actions.copy()
+        check_attribute_actions.set_index_keys(*list(TEST_KEY.keys()))
+        check_many_encrypted_items(
+            actual=encrypted_items,
+            expected=decrypted_items,
+            attribute_actions=check_attribute_actions,
+            transformer=ddb_to_dict,
+        )
 
-    _cleanup_items(encrypted=e_client, write_transformer=dict_to_ddb, table_name=table_name)
+    finally:
+        _cleanup_items(encrypted=e_client, write_transformer=dict_to_ddb, table_name=table_name)
 
     raw_scan_result = client.scan(TableName=table_name, ConsistentRead=True)
     e_scan_result = e_client.scan(TableName=table_name, ConsistentRead=True)

test/integration/encrypted/test_client.py

@@ -54,3 +54,19 @@ def test_ephemeral_batch_item_cycle_kms(ddb_table_name, aws_kms_cmp, parametrize
     functional_test_utils.client_cycle_batch_items_check(
         aws_kms_cmp, parametrized_actions, parametrized_item, ddb_table_name
     )
+
+
+def test_ephemeral_batch_item_cycle_scan_paginator(ddb_table_name, some_cmps, parametrized_actions, parametrized_item):
+    """Test a small number of curated CMPs against a small number of curated items using the scan paginator."""
+    functional_test_utils.client_cycle_batch_items_check_scan_paginator(
+        some_cmps, parametrized_actions, parametrized_item, ddb_table_name
+    )
+
+
+def test_ephemeral_batch_item_cycle_scan_paginator_kms(
+    ddb_table_name, aws_kms_cmp, parametrized_actions, parametrized_item
+):
+    """Test a the AWS KMS CMP against a small number of curated items using the scan paginator."""
+    functional_test_utils.client_cycle_batch_items_check_scan_paginator(
+        aws_kms_cmp, parametrized_actions, parametrized_item, ddb_table_name
+    )