all message format versions

Author: lucasmcdonald3

src/aws_encryption_sdk/internal/formatting/serialize.py

@@ -189,7 +189,13 @@ def serialize_header(header, signer=None):
         raise SerializationError("Unrecognized message format version: {}".format(header.version))
 
 
-def _serialize_header_auth_v1(algorithm, header, data_encryption_key, signer=None):
+def _serialize_header_auth_v1(
+    algorithm,
+    header,
+    data_encryption_key,
+    signer=None,
+    required_ec_bytes=None
+):
     """Creates serialized header authentication data for messages in serialization version V1.
 
     :param algorithm: Algorithm to use for encryption
@@ -198,16 +204,35 @@ def _serialize_header_auth_v1(algorithm, header, data_encryption_key, signer=Non
     :param bytes data_encryption_key: Data key with which to encrypt message
     :param signer: Cryptographic signer object (optional)
     :type signer: aws_encryption_sdk.Signer
+    :param required_encryption_context_bytes: Serialized encryption context items
+        for all items whose keys are in the required_encryption_context list.
+        This is ONLY processed if using the aws-cryptographic-materialproviders library
+        AND its required encryption context CMM. (optional)
+    :type required_encryption_context_bytes: bytes
     :returns: Serialized header authentication data
     :rtype: bytes
     """
-    header_auth = encrypt(
-        algorithm=algorithm,
-        key=data_encryption_key,
-        plaintext=b"",
-        associated_data=header,
-        iv=header_auth_iv(algorithm),
-    )
+    if required_ec_bytes is None:
+        header_auth = encrypt(
+            algorithm=algorithm,
+            key=data_encryption_key,
+            plaintext=b"",
+            associated_data=header,
+            iv=header_auth_iv(algorithm),
+        )
+    else:
+        header_auth = encrypt(
+            algorithm=algorithm,
+            key=data_encryption_key,
+            plaintext=b"",
+            # The AAD MUST be the concatenation of the serialized message header body and the serialization
+            # of encryption context to only authenticate. The encryption context to only authenticate MUST
+            # be the encryption context in the encryption materials filtered to only contain key value
+            # pairs listed in the encryption material's required encryption context keys serialized
+            # according to the encryption context serialization specification.
+            associated_data=header + required_ec_bytes,
+            iv=header_auth_iv(algorithm),
+        )
     output = struct.pack(
         ">{iv_len}s{tag_len}s".format(iv_len=algorithm.iv_len, tag_len=algorithm.tag_len),
         header_auth.iv,
@@ -292,8 +317,7 @@ def serialize_header_auth(
     :param required_encryption_context_bytes: Serialized encryption context items
         for all items whose keys are in the required_encryption_context list.
         This is ONLY processed if using the aws-cryptographic-materialproviders library
-        AND its required encryption context CMM
-        AND if using the v2 message format. (optional)
+        AND its required encryption context CMM. (optional)
     :type required_encryption_context_bytes: bytes
     :returns: Serialized header authentication data
     :rtype: bytes