add keyring trace and integrate into updated encrytion/decryption materials

Author: mattsb42-aws

src/aws_encryption_sdk/identifiers.py

@@ -328,3 +328,13 @@ class ContentAADString(Enum):
     FRAME_STRING_ID = b"AWSKMSEncryptionClient Frame"
     FINAL_FRAME_STRING_ID = b"AWSKMSEncryptionClient Final Frame"
     NON_FRAMED_STRING_ID = b"AWSKMSEncryptionClient Single Block"
+
+
+class KeyRingTraceFlag(Enum):
+    """KeyRing Trace actions."""
+
+    WRAPPING_KEY_GENERATED_DATA_KEY = 1
+    WRAPPING_KEY_ENCRYPTED_DATA_KEY = 2
+    WRAPPING_KEY_DECRYPTED_DATA_KEY = 3
+    WRAPPING_KEY_SIGNED_ENC_CTX = 4
+    WRAPPING_KEY_VERIFIED_ENC_CTX = 5

src/aws_encryption_sdk/materials_managers/__init__.py

@@ -16,10 +16,11 @@
 """
 import attr
 import six
+from attr.validators import deep_iterable, deep_mapping, instance_of, optional
 
-from ..identifiers import Algorithm
+from ..identifiers import Algorithm, KeyRingTraceFlag
 from ..internal.utils.streams import ROStream
-from ..structures import DataKey
+from ..structures import DataKey, EncryptedDataKey, KeyRingTrace
 
 
 @attr.s(hash=False)
@@ -51,28 +52,87 @@ class EncryptionMaterialsRequest(object):
     )
 
 
-@attr.s(hash=False)
-class EncryptionMaterials(object):
+@attr.s
+class CryptographicMaterials(object):
+    """Cryptographic materials core.
+
+    .. versionadded:: 1.5.0
+
+    :param Algorithm algorithm: Algorithm to use for encrypting message
+    :param dict encryption_context: Encryption context tied to `encrypted_data_keys`
+    :param DataKey data_encryption_key: Plaintext data key to use for encrypting message
+    :param encrypted_data_keys: List of encrypted data keys
+    :type encrypted_data_keys: list of :class:`EncryptedDataKey`
+    :param keyring_trace: Any KeyRing trace entries
+    :type keyring_trace: list of :class:`KeyRingTrace`
+    """
+
+    algorithm = attr.ib(validator=optional(instance_of(Algorithm)))
+    encryption_context = attr.ib(
+        validator=optional(
+            deep_mapping(key_validator=instance_of(six.string_types), value_validator=instance_of(six.string_types))
+        )
+    )
+    data_encryption_key = attr.ib(default=None, validator=optional(instance_of(DataKey)))
+    encrypted_data_keys = attr.ib(
+        default=attr.Factory(list), validator=optional(deep_iterable(member_validator=instance_of(EncryptedDataKey)))
+    )
+    keyring_trace = attr.ib(
+        default=attr.Factory(list), validator=optional(deep_iterable(member_validator=instance_of(KeyRingTrace)))
+    )
+
+
+@attr.s(hash=False, init=False)
+class EncryptionMaterials(CryptographicMaterials):
     """Encryption materials returned by a crypto material manager's `get_encryption_materials` method.
 
     .. versionadded:: 1.3.0
 
-    :param algorithm: Algorithm to use for encrypting message
-    :type algorithm: aws_encryption_sdk.identifiers.Algorithm
-    :param data_encryption_key: Plaintext data key to use for encrypting message
-    :type data_encryption_key: aws_encryption_sdk.structures.DataKey
-    :param encrypted_data_keys: List of encrypted data keys
-    :type encrypted_data_keys: list of `aws_encryption_sdk.structures.EncryptedDataKey`
+    .. versionadded:: 1.5.0
+
+        The **keyring_trace** parameter.
+
+    .. versionadded:: 1.5.0
+
+        Most parameters are now optional.
+
+    :param Algorithm algorithm: Algorithm to use for encrypting message
+    :param DataKey data_encryption_key: Plaintext data key to use for encrypting message (optional)
+    :param encrypted_data_keys: List of encrypted data keys (optional)
+    :type encrypted_data_keys: list of :class:`EncryptedDataKey`
     :param dict encryption_context: Encryption context tied to `encrypted_data_keys`
-    :param bytes signing_key: Encoded signing key
+    :param bytes signing_key: Encoded signing key (optional)
+    :param keyring_trace: Any KeyRing trace entries (optional)
+    :type keyring_trace: list of :class:`KeyRingTrace`
     """
 
-    algorithm = attr.ib(validator=attr.validators.instance_of(Algorithm))
-    data_encryption_key = attr.ib(validator=attr.validators.instance_of(DataKey))
-    encrypted_data_keys = attr.ib(validator=attr.validators.instance_of(set))
-    encryption_context = attr.ib(validator=attr.validators.instance_of(dict))
     signing_key = attr.ib(default=None, validator=attr.validators.optional(attr.validators.instance_of(bytes)))
 
+    def __init__(
+        self,
+        algorithm=None,
+        data_encryption_key=None,
+        encrypted_data_keys=None,
+        encryption_context=None,
+        signing_key=None,
+        **kwargs
+    ):
+        if algorithm is None:
+            raise TypeError("algorithm must not be None")
+
+        if encryption_context is None:
+            raise TypeError("encryption_context must not be None")
+
+        super(EncryptionMaterials, self).__init__(
+            algorithm=algorithm,
+            encryption_context=encryption_context,
+            data_encryption_key=data_encryption_key,
+            encrypted_data_keys=encrypted_data_keys,
+            **kwargs
+        )
+        self.signing_key = signing_key
+        attr.validate(self)
+
 
 @attr.s(hash=False)
 class DecryptionMaterialsRequest(object):
@@ -92,16 +152,64 @@ class DecryptionMaterialsRequest(object):
     encryption_context = attr.ib(validator=attr.validators.instance_of(dict))
 
 
-@attr.s(hash=False)
-class DecryptionMaterials(object):
+_DEFAULT_SENTINEL = object()
+
+
+@attr.s(hash=False, init=False)
+class DecryptionMaterials(CryptographicMaterials):
     """Decryption materials returned by a crypto material manager's `decrypt_materials` method.
 
     .. versionadded:: 1.3.0
 
-    :param data_key: Plaintext data key to use with message decryption
-    :type data_key: aws_encryption_sdk.structures.DataKey
-    :param bytes verification_key: Raw signature verification key
+    .. versionadded:: 1.5.0
+
+        The **algorithm**, **data_encryption_key**, **encrypted_data_keys**,
+        **encryption_context**, and **keyring_trace** parameters.
+
+    .. versionadded:: 1.5.0
+
+        All parameters are now optional.
+
+    :param Algorithm algorithm: Algorithm to use for encrypting message (optional)
+    :param DataKey data_encryption_key: Plaintext data key to use for encrypting message (optional)
+    :param encrypted_data_keys: List of encrypted data keys (optional)
+    :type encrypted_data_keys: list of :class:`EncryptedDataKey`
+    :param dict encryption_context: Encryption context tied to `encrypted_data_keys` (optional)
+    :param bytes verification_key: Raw signature verification key (optional)
+    :param keyring_trace: Any KeyRing trace entries (optional)
+    :type keyring_trace: list of :class:`KeyRingTrace`
     """
 
-    data_key = attr.ib(validator=attr.validators.instance_of(DataKey))
     verification_key = attr.ib(default=None, validator=attr.validators.optional(attr.validators.instance_of(bytes)))
+
+    def __init__(self, data_key=_DEFAULT_SENTINEL, verification_key=None, **kwargs):
+        if any(
+            (
+                data_key is _DEFAULT_SENTINEL and "data_encryption_key" not in kwargs,
+                data_key is not _DEFAULT_SENTINEL and "data_encryption_key" in kwargs,
+            )
+        ):
+            raise TypeError("Exactly one of data_key or data_encryption_key must be set")
+
+        if data_key is not _DEFAULT_SENTINEL and "data_encryption_key" not in kwargs:
+            kwargs["data_encryption_key"] = data_key
+
+        for legacy_missing in ("algorithm", "encryption_context"):
+            if legacy_missing not in kwargs:
+                kwargs[legacy_missing] = None
+
+        super(DecryptionMaterials, self).__init__(**kwargs)
+
+        self.verification_key = verification_key
+        attr.validate(self)
+
+    @property
+    def data_key(self):
+        """Backwards-compatible shim."""
+        return self.data_encryption_key
+
+    @data_key.setter
+    def data_key(self, value):
+        # type: (DataKey) -> None
+        """Backwards-compatible shim."""
+        self.data_encryption_key = value

src/aws_encryption_sdk/structures.py

@@ -13,6 +13,7 @@
 """Public data structures for aws_encryption_sdk."""
 import attr
 import six
+from attr.validators import deep_iterable, instance_of
 
 import aws_encryption_sdk.identifiers
 from aws_encryption_sdk.internal.str_ops import to_bytes, to_str
@@ -104,3 +105,18 @@ class EncryptedDataKey(object):
 
     key_provider = attr.ib(hash=True, validator=attr.validators.instance_of(MasterKeyInfo))
     encrypted_data_key = attr.ib(hash=True, validator=attr.validators.instance_of(bytes))
+
+
+@attr.s
+class KeyRingTrace(object):
+    """Record of all actions that a KeyRing performed with a wrapping key.
+
+    :param MasterKeyInfo wrapping_key: Wrapping key used
+    :param flags: Actions performed
+    :type flags: set of :class:`KeyRingTraceFlag`
+    """
+
+    wrapping_key = attr.ib(validator=instance_of(MasterKeyInfo))
+    flags = attr.ib(
+        validator=deep_iterable(member_validator=instance_of(aws_encryption_sdk.identifiers.KeyRingTraceFlag))
+    )

test/unit/test_material_managers.py

@@ -12,66 +12,93 @@
 # language governing permissions and limitations under the License.
 """Test suite for aws_encryption_sdk.materials_managers"""
 import pytest
-from mock import MagicMock
+from mock import MagicMock, sentinel
 from pytest_mock import mocker  # noqa pylint: disable=unused-import
 
-from aws_encryption_sdk.identifiers import Algorithm
+from aws_encryption_sdk.identifiers import KeyRingTraceFlag
+from aws_encryption_sdk.internal.defaults import ALGORITHM
 from aws_encryption_sdk.internal.utils.streams import ROStream
 from aws_encryption_sdk.materials_managers import (
+    CryptographicMaterials,
     DecryptionMaterials,
     DecryptionMaterialsRequest,
     EncryptionMaterials,
     EncryptionMaterialsRequest,
 )
-from aws_encryption_sdk.structures import DataKey
+from aws_encryption_sdk.structures import DataKey, KeyRingTrace, MasterKeyInfo
 
 pytestmark = [pytest.mark.unit, pytest.mark.local]
 
+_DATA_KEY = DataKey(
+    key_provider=MasterKeyInfo(provider_id="Provider", key_info=b"Info"),
+    data_key=b"1234567890123456789012",
+    encrypted_data_key=b"asdf",
+)
 
 _VALID_KWARGS = {
+    "CryptographicMaterials": dict(
+        algorithm=ALGORITHM,
+        encryption_context={"additional": "data"},
+        data_encryption_key=_DATA_KEY,
+        encrypted_data_keys=[],
+        keyring_trace=[
+            KeyRingTrace(
+                wrapping_key=MasterKeyInfo(provider_id="Provider", key_info=b"Info"),
+                flags={KeyRingTraceFlag.WRAPPING_KEY_GENERATED_DATA_KEY},
+            )
+        ],
+    ),
     "EncryptionMaterialsRequest": dict(
         encryption_context={},
         plaintext_rostream=MagicMock(__class__=ROStream),
         frame_length=5,
-        algorithm=MagicMock(__class__=Algorithm),
+        algorithm=ALGORITHM,
         plaintext_length=5,
     ),
     "EncryptionMaterials": dict(
-        algorithm=MagicMock(__class__=Algorithm),
-        data_encryption_key=MagicMock(__class__=DataKey),
+        algorithm=ALGORITHM,
+        data_encryption_key=_DATA_KEY,
         encrypted_data_keys=set([]),
         encryption_context={},
         signing_key=b"",
     ),
-    "DecryptionMaterialsRequest": dict(
-        algorithm=MagicMock(__class__=Algorithm), encrypted_data_keys=set([]), encryption_context={}
+    "DecryptionMaterialsRequest": dict(algorithm=ALGORITHM, encrypted_data_keys=set([]), encryption_context={}),
+    "DecryptionMaterials": dict(
+        data_key=_DATA_KEY, verification_key=b"ex_verification_key", algorithm=ALGORITHM, encryption_context={}
     ),
-    "DecryptionMaterials": dict(data_key=MagicMock(__class__=DataKey), verification_key=b"ex_verification_key"),
 }
+_REMOVE = object()
 
 
 @pytest.mark.parametrize(
     "attr_class, invalid_kwargs",
     (
+        (CryptographicMaterials, dict(algorithm=1234)),
+        (CryptographicMaterials, dict(encryption_context=1234)),
+        (CryptographicMaterials, dict(data_encryption_key=1234)),
+        (CryptographicMaterials, dict(encrypted_data_keys=1234)),
+        (CryptographicMaterials, dict(keyring_trace=1234)),
         (EncryptionMaterialsRequest, dict(encryption_context=None)),
         (EncryptionMaterialsRequest, dict(frame_length="not an int")),
         (EncryptionMaterialsRequest, dict(algorithm="not an Algorithm or None")),
         (EncryptionMaterialsRequest, dict(plaintext_length="not an int or None")),
         (EncryptionMaterials, dict(algorithm=None)),
-        (EncryptionMaterials, dict(data_encryption_key=None)),
-        (EncryptionMaterials, dict(encrypted_data_keys=None)),
         (EncryptionMaterials, dict(encryption_context=None)),
         (EncryptionMaterials, dict(signing_key=u"not bytes or None")),
         (DecryptionMaterialsRequest, dict(algorithm=None)),
         (DecryptionMaterialsRequest, dict(encrypted_data_keys=None)),
         (DecryptionMaterialsRequest, dict(encryption_context=None)),
-        (DecryptionMaterials, dict(data_key=None)),
         (DecryptionMaterials, dict(verification_key=5555)),
+        (DecryptionMaterials, dict(data_key=_DATA_KEY, data_encryption_key=_DATA_KEY)),
+        (DecryptionMaterials, dict(data_key=_REMOVE, data_encryption_key=_REMOVE)),
     ),
 )
 def test_attributes_fails(attr_class, invalid_kwargs):
     kwargs = _VALID_KWARGS[attr_class.__name__].copy()
     kwargs.update(invalid_kwargs)
+    purge_keys = [key for key, val in kwargs.items() if val is _REMOVE]
+    for key in purge_keys:
+        del kwargs[key]
     with pytest.raises(TypeError):
         attr_class(**kwargs)
 
@@ -85,14 +112,29 @@ def test_encryption_materials_request_attributes_defaults():
 
 def test_encryption_materials_defaults():
     test = EncryptionMaterials(
-        algorithm=MagicMock(__class__=Algorithm),
-        data_encryption_key=MagicMock(__class__=DataKey),
-        encrypted_data_keys=set([]),
-        encryption_context={},
+        algorithm=ALGORITHM, data_encryption_key=_DATA_KEY, encrypted_data_keys=set([]), encryption_context={}
     )
     assert test.signing_key is None
 
 
 def test_decryption_materials_defaults():
-    test = DecryptionMaterials(data_key=MagicMock(__class__=DataKey))
+    test = DecryptionMaterials(data_key=_DATA_KEY)
     assert test.verification_key is None
+    assert test.algorithm is None
+    assert test.encryption_context is None
+
+
+def test_decryption_materials_legacy_data_key_get():
+    test = DecryptionMaterials(data_encryption_key=_DATA_KEY)
+
+    assert test.data_encryption_key is _DATA_KEY
+    assert test.data_key is _DATA_KEY
+
+
+def test_decryption_materials_legacy_data_key_set():
+    test = DecryptionMaterials(data_encryption_key=_DATA_KEY)
+
+    test.data_key = sentinel.data_key
+
+    assert test.data_encryption_key is sentinel.data_key
+    assert test.data_key is sentinel.data_key

test/unit/test_material_managers_default.py

@@ -22,7 +22,7 @@
 from aws_encryption_sdk.key_providers.base import MasterKeyProvider
 from aws_encryption_sdk.materials_managers import EncryptionMaterials
 from aws_encryption_sdk.materials_managers.default import DefaultCryptoMaterialsManager
-from aws_encryption_sdk.structures import DataKey
+from aws_encryption_sdk.structures import DataKey, EncryptedDataKey
 
 pytestmark = [pytest.mark.unit, pytest.mark.local]
 
@@ -34,7 +34,7 @@ def patch_for_dcmm_encrypt(mocker):
     DefaultCryptoMaterialsManager._generate_signing_key_and_update_encryption_context.return_value = mock_signing_key
     mocker.patch.object(aws_encryption_sdk.materials_managers.default, "prepare_data_keys")
     mock_data_encryption_key = MagicMock(__class__=DataKey)
-    mock_encrypted_data_keys = set([mock_data_encryption_key])
+    mock_encrypted_data_keys = set([MagicMock(__class__=EncryptedDataKey)])
     result_pair = mock_data_encryption_key, mock_encrypted_data_keys
     aws_encryption_sdk.materials_managers.default.prepare_data_keys.return_value = result_pair
     yield result_pair, mock_signing_key
@@ -128,7 +128,7 @@ def test_get_encryption_materials(patch_for_dcmm_encrypt):
     assert isinstance(test, EncryptionMaterials)
     assert test.algorithm is cmm.algorithm
     assert test.data_encryption_key is patch_for_dcmm_encrypt[0][0]
-    assert test.encrypted_data_keys is patch_for_dcmm_encrypt[0][1]
+    assert test.encrypted_data_keys == patch_for_dcmm_encrypt[0][1]
     assert test.encryption_context == encryption_context
     assert test.signing_key == patch_for_dcmm_encrypt[1]