fix AWS KMS encryption context calculation to correctly handle binary values (#45)

* fix AWS KMS encryption context calculation to correctly handle binary values

Author: mattsb42-aws

src/dynamodb_encryption_sdk/material_providers/aws_kms.py

@@ -301,7 +301,7 @@ def _attribute_to_value(self, attribute):
         """
         attribute_type, attribute_value = list(attribute.items())[0]
         if attribute_type == 'B':
-            return base64.b64encode(attribute_value.value).decode(TEXT_ENCODING)
+            return base64.b64encode(attribute_value).decode(TEXT_ENCODING)
         if attribute_type in ('S', 'N'):
             return attribute_value
         raise ValueError('Attribute of type "{}" cannot be used in KMS encryption context.'.format(attribute_type))

test/integration/material_providers/test_aws_kms.py

@@ -12,18 +12,22 @@
 # language governing permissions and limitations under the License.
 """Integration tests for ``dynamodb_encryption_sdk.material_providers.aws_kms``."""
 import logging
+import itertools
 
 import hypothesis
 import pytest
 
 from dynamodb_encryption_sdk.encrypted import CryptoConfig
-from dynamodb_encryption_sdk.identifiers import USER_AGENT_SUFFIX
-from dynamodb_encryption_sdk.structures import EncryptionContext
+from dynamodb_encryption_sdk.identifiers import CryptoAction, USER_AGENT_SUFFIX
+from dynamodb_encryption_sdk.structures import AttributeActions, EncryptionContext
+from dynamodb_encryption_sdk.transform import dict_to_ddb
 from ..integration_test_utils import aws_kms_cmp  # noqa pylint: disable=unused-import
 from ..integration_test_utils import functional_test_utils, hypothesis_strategies
 
 pytestmark = pytest.mark.integ
 
+_primary_key_names = ('partition_key', 'sort_key')
+
 
 def pytest_generate_tests(metafunc):
     functional_test_utils.set_parametrized_actions(metafunc)
@@ -38,6 +42,35 @@ def test_verify_user_agent(aws_kms_cmp, caplog):
     assert USER_AGENT_SUFFIX in caplog.text
 
 
+def _many_items():
+    values = ('a string', 1234, b'binary \x00\x88 value')
+    partition_keys = (('partition_key', value) for value in values)
+    sort_keys = (('sort_key', value) for value in values)
+    for pairs in itertools.product(partition_keys, sort_keys):
+        item = dict(pairs)
+        yield pytest.param(item, id=str(item))
+
+
+@pytest.mark.parametrize('item', _many_items())
+def test_aws_kms_diverse_indexes(aws_kms_cmp, item):
+    """Verify that AWS KMS cycle works for items with all possible combinations for primary index attribute types."""
+    crypto_config = CryptoConfig(
+        materials_provider=aws_kms_cmp,
+        encryption_context=EncryptionContext(
+            partition_key_name='partition_key',
+            sort_key_name='sort_key',
+            attributes=dict_to_ddb(item)
+        ),
+        attribute_actions=AttributeActions(
+            attribute_actions={
+                key: CryptoAction.SIGN_ONLY
+                for key in _primary_key_names
+            }
+        )
+    )
+    functional_test_utils.cycle_item_check(item, crypto_config)
+
+
 def test_aws_kms_item_cycle(aws_kms_cmp, parametrized_actions, parametrized_item):
     crypto_config = CryptoConfig(
         materials_provider=aws_kms_cmp,

test/unit/material_providers/test_aws_kms.py

@@ -318,7 +318,7 @@ def test_select_id(default_kms_cmp):
 
 # TODO: vectorize
 @pytest.mark.parametrize('attribute, expected_value', (
-    ({'B': Binary(b'\x00\x01\x02\x03')}, 'AAECAw=='),
+    ({'B': b'\x00\x01\x02\x03'}, 'AAECAw=='),
     ({'S': 'some string value'}, 'some string value'),
     ({'N': '55.2'}, '55.2')
 ))