feat: immutable cryptographic materials (for keyrings) (aws#231)

* feat: change keyring materials modifiers to immutable constructors

* feat: add key_name parameter and key_namespace accessor to MasterKeyInfo

* feat: adopt immutable materials for KMS keyring

* feat: adopt immutable materials for multi-keyring

* fix: revise test helpers to use key IDs that are not used in tests

* feat: add key_name parameter and key_namespace accessor to MasterKeyInfo pt2

* feat: fix raw keyrings and move to immutable materials

* adopt immutable materials
* add encryption context entries to keyring trace awslabs/aws-encryption-sdk-specification#32
* use the key name, not the provider info, in the keyring trace
* fix the keyring trace tests to accurately fail when trace is wrong

* fix: fix incorrectly placed test assertions

* chore: clean up linting

* docs: update src/aws_encryption_sdk/keyrings/raw.py

Co-Authored-By: Wesley Rosenblum <55108558+WesleyRosenblum@users.noreply.github.com>

* chore: remove unnecessary copies on encrypt and avoid overwriting pointers on decrypt

* chore: finish copy cleanup

Co-authored-by: Wesley Rosenblum <55108558+WesleyRosenblum@users.noreply.github.com>

Author: mattsb42-aws

src/aws_encryption_sdk/internal/formatting/serialize.py

@@ -316,6 +316,6 @@ def serialize_wrapped_key(key_provider, wrapping_algorithm, wrapping_key_id, enc
         )
         key_ciphertext = encrypted_wrapped_key.ciphertext + encrypted_wrapped_key.tag
     return EncryptedDataKey(
-        key_provider=MasterKeyInfo(provider_id=key_provider.provider_id, key_info=key_info),
+        key_provider=MasterKeyInfo(provider_id=key_provider.provider_id, key_info=key_info, key_name=wrapping_key_id),
         encrypted_data_key=key_ciphertext,
     )

src/aws_encryption_sdk/keyrings/aws_kms/__init__.py

@@ -180,25 +180,26 @@ class _AwsKmsSingleCmkKeyring(Keyring):
     def on_encrypt(self, encryption_materials):
         # type: (EncryptionMaterials) -> EncryptionMaterials
         trace_info = MasterKeyInfo(provider_id=_PROVIDER_ID, key_info=self._key_id)
+        new_materials = encryption_materials
         try:
-            if encryption_materials.data_encryption_key is None:
+            if new_materials.data_encryption_key is None:
                 plaintext_key, encrypted_key = _do_aws_kms_generate_data_key(
                     client_supplier=self._client_supplier,
                     key_name=self._key_id,
-                    encryption_context=encryption_materials.encryption_context,
-                    algorithm=encryption_materials.algorithm,
+                    encryption_context=new_materials.encryption_context,
+                    algorithm=new_materials.algorithm,
                     grant_tokens=self._grant_tokens,
                 )
-                encryption_materials.add_data_encryption_key(
+                new_materials = new_materials.with_data_encryption_key(
                     data_encryption_key=plaintext_key,
                     keyring_trace=KeyringTrace(wrapping_key=trace_info, flags=_GENERATE_FLAGS),
                 )
             else:
                 encrypted_key = _do_aws_kms_encrypt(
                     client_supplier=self._client_supplier,
                     key_name=self._key_id,
-                    plaintext_data_key=encryption_materials.data_encryption_key,
-                    encryption_context=encryption_materials.encryption_context,
+                    plaintext_data_key=new_materials.data_encryption_key,
+                    encryption_context=new_materials.encryption_context,
                     grant_tokens=self._grant_tokens,
                 )
         except Exception:  # pylint: disable=broad-except
@@ -207,30 +208,30 @@ def on_encrypt(self, encryption_materials):
             _LOGGER.exception(message)
             raise EncryptKeyError(message)
 
-        encryption_materials.add_encrypted_data_key(
+        return new_materials.with_encrypted_data_key(
             encrypted_data_key=encrypted_key, keyring_trace=KeyringTrace(wrapping_key=trace_info, flags=_ENCRYPT_FLAGS)
         )
 
-        return encryption_materials
-
     def on_decrypt(self, decryption_materials, encrypted_data_keys):
         # type: (DecryptionMaterials, Iterable[EncryptedDataKey]) -> DecryptionMaterials
+        new_materials = decryption_materials
+
         for edk in encrypted_data_keys:
-            if decryption_materials.data_encryption_key is not None:
-                return decryption_materials
+            if new_materials.data_encryption_key is not None:
+                return new_materials
 
             if (
                 edk.key_provider.provider_id == _PROVIDER_ID
                 and edk.key_provider.key_info.decode("utf-8") == self._key_id
             ):
-                decryption_materials = _try_aws_kms_decrypt(
+                new_materials = _try_aws_kms_decrypt(
                     client_supplier=self._client_supplier,
-                    decryption_materials=decryption_materials,
+                    decryption_materials=new_materials,
                     grant_tokens=self._grant_tokens,
                     encrypted_data_key=edk,
                 )
 
-        return decryption_materials
+        return new_materials
 
 
 @attr.s
@@ -258,19 +259,21 @@ def on_encrypt(self, encryption_materials):
 
     def on_decrypt(self, decryption_materials, encrypted_data_keys):
         # type: (DecryptionMaterials, Iterable[EncryptedDataKey]) -> DecryptionMaterials
+        new_materials = decryption_materials
+
         for edk in encrypted_data_keys:
-            if decryption_materials.data_encryption_key is not None:
-                return decryption_materials
+            if new_materials.data_encryption_key is not None:
+                return new_materials
 
             if edk.key_provider.provider_id == _PROVIDER_ID:
-                decryption_materials = _try_aws_kms_decrypt(
+                new_materials = _try_aws_kms_decrypt(
                     client_supplier=self._client_supplier,
-                    decryption_materials=decryption_materials,
+                    decryption_materials=new_materials,
                     grant_tokens=self._grant_tokens,
                     encrypted_data_key=edk,
                 )
 
-        return decryption_materials
+        return new_materials
 
 
 def _try_aws_kms_decrypt(client_supplier, decryption_materials, grant_tokens, encrypted_data_key):
@@ -293,14 +296,12 @@ def _try_aws_kms_decrypt(client_supplier, decryption_materials, grant_tokens, en
     except Exception:  # pylint: disable=broad-except
         # We intentionally WANT to catch all exceptions here
         _LOGGER.exception("Unable to decrypt encrypted data key from %s", encrypted_data_key.key_provider)
-    else:
-        decryption_materials.add_data_encryption_key(
-            data_encryption_key=plaintext_key,
-            keyring_trace=KeyringTrace(wrapping_key=encrypted_data_key.key_provider, flags=_DECRYPT_FLAGS),
-        )
         return decryption_materials
 
-    return decryption_materials
+    return decryption_materials.with_data_encryption_key(
+        data_encryption_key=plaintext_key,
+        keyring_trace=KeyringTrace(wrapping_key=encrypted_data_key.key_provider, flags=_DECRYPT_FLAGS),
+    )
 
 
 def _do_aws_kms_decrypt(client_supplier, key_name, encrypted_data_key, encryption_context, grant_tokens):

src/aws_encryption_sdk/keyrings/multi.py

@@ -67,20 +67,21 @@ def on_encrypt(self, encryption_materials):
                 "and encryption materials do not already contain a plaintext data key."
             )
 
+        new_materials = encryption_materials
+
         # Call on_encrypt on the generator keyring if it is provided
         if self.generator is not None:
-
-            encryption_materials = self.generator.on_encrypt(encryption_materials=encryption_materials)
+            new_materials = self.generator.on_encrypt(encryption_materials=new_materials)
 
         # Check if data key is generated
-        if encryption_materials.data_encryption_key is None:
+        if new_materials.data_encryption_key is None:
             raise GenerateKeyError("Unable to generate data encryption key.")
 
         # Call on_encrypt on all other keyrings
         for keyring in self.children:
-            encryption_materials = keyring.on_encrypt(encryption_materials=encryption_materials)
+            new_materials = keyring.on_encrypt(encryption_materials=new_materials)
 
-        return encryption_materials
+        return new_materials
 
     def on_decrypt(self, decryption_materials, encrypted_data_keys):
         # type: (DecryptionMaterials, Iterable[EncryptedDataKey]) -> DecryptionMaterials
@@ -92,10 +93,13 @@ def on_decrypt(self, decryption_materials, encrypted_data_keys):
         :rtype: DecryptionMaterials
         """
         # Call on_decrypt on all keyrings till decryption is successful
+        new_materials = decryption_materials
         for keyring in self._decryption_keyrings:
-            if decryption_materials.data_encryption_key is not None:
-                return decryption_materials
-            decryption_materials = keyring.on_decrypt(
-                decryption_materials=decryption_materials, encrypted_data_keys=encrypted_data_keys
+            if new_materials.data_encryption_key is not None:
+                return new_materials
+
+            new_materials = keyring.on_decrypt(
+                decryption_materials=new_materials, encrypted_data_keys=encrypted_data_keys
             )
-        return decryption_materials
+
+        return new_materials