feat: remove wrapping_algorithm input parameter from RawAESKeyring (aws#247)

* feat: remove wrapping_algorithm input parameter from RawAESKeyring

* docs: make RawAESKeyring key material length error message deterministic and check for exact allowed options

* docs: tweak wording explaining wrapping key generation

Author: mattsb42-aws

examples/src/keyring/raw_aes/raw_aes.py

@@ -10,7 +10,6 @@
 import os
 
 import aws_encryption_sdk
-from aws_encryption_sdk.identifiers import WrappingAlgorithm
 from aws_encryption_sdk.keyrings.raw import RawAESKeyring
 
 
@@ -30,14 +29,10 @@ def run(source_plaintext):
         "the data you are handling": "is what you think it is",
     }
 
-    # Choose the wrapping algorithm for the keyring to use.
-    wrapping_algorithm = WrappingAlgorithm.AES_256_GCM_IV12_TAG16_NO_PADDING
-
-    # Generate an AES key to use with your keyring.
-    # The key size depends on the wrapping algorithm.
+    # Generate a 256-bit (32 byte) AES key to use with your keyring.
     #
     # In practice, you should get this key from a secure key management system such as an HSM.
-    key = os.urandom(wrapping_algorithm.algorithm.kdf_input_len)
+    key = os.urandom(32)
 
     # Create the keyring that determines how your data keys are protected.
     keyring = RawAESKeyring(
@@ -50,7 +45,6 @@ def run(source_plaintext):
         key_namespace="some managed raw keys",
         key_name=b"my AES wrapping key",
         wrapping_key=key,
-        wrapping_algorithm=wrapping_algorithm,
     )
 
     # Encrypt your plaintext data.

src/aws_encryption_sdk/keyrings/raw.py

@@ -6,7 +6,7 @@
 
 import attr
 import six
-from attr.validators import instance_of, optional
+from attr.validators import in_, instance_of, optional
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey
@@ -76,7 +76,6 @@ class RawAESKeyring(Keyring):
     :param str key_namespace: String defining the keyring.
     :param bytes key_name: Key ID
     :param bytes wrapping_key: Encryption key with which to wrap plaintext data key.
-    :param WrappingAlgorithm wrapping_algorithm: Wrapping Algorithm with which to wrap plaintext data key.
 
     .. note::
 
@@ -86,11 +85,28 @@ class RawAESKeyring(Keyring):
     key_namespace = attr.ib(validator=instance_of(six.string_types))
     key_name = attr.ib(validator=instance_of(six.binary_type))
     _wrapping_key = attr.ib(repr=False, validator=instance_of(six.binary_type))
-    _wrapping_algorithm = attr.ib(repr=False, validator=instance_of(WrappingAlgorithm))
 
     def __attrs_post_init__(self):
         # type: () -> None
         """Prepares initial values not handled by attrs."""
+        key_size_to_wrapping_algorithm = {
+            wrapper.algorithm.kdf_input_len: wrapper
+            for wrapper in (
+                WrappingAlgorithm.AES_128_GCM_IV12_TAG16_NO_PADDING,
+                WrappingAlgorithm.AES_192_GCM_IV12_TAG16_NO_PADDING,
+                WrappingAlgorithm.AES_256_GCM_IV12_TAG16_NO_PADDING,
+            )
+        }
+
+        try:
+            self._wrapping_algorithm = key_size_to_wrapping_algorithm[len(self._wrapping_key)]
+        except KeyError:
+            raise ValueError(
+                "Invalid wrapping key length. Must be one of {} bytes.".format(
+                    sorted(key_size_to_wrapping_algorithm.keys())
+                )
+            )
+
         self._key_provider = MasterKeyInfo(provider_id=self.key_namespace, key_info=self.key_name)
 
         self._wrapping_key_structure = WrappingKey(
@@ -244,7 +260,18 @@ class RawRSAKeyring(Keyring):
 
     key_namespace = attr.ib(validator=instance_of(six.string_types))
     key_name = attr.ib(validator=instance_of(six.binary_type))
-    _wrapping_algorithm = attr.ib(repr=False, validator=instance_of(WrappingAlgorithm))
+    _wrapping_algorithm = attr.ib(
+        repr=False,
+        validator=in_(
+            (
+                WrappingAlgorithm.RSA_PKCS1,
+                WrappingAlgorithm.RSA_OAEP_SHA1_MGF1,
+                WrappingAlgorithm.RSA_OAEP_SHA256_MGF1,
+                WrappingAlgorithm.RSA_OAEP_SHA384_MGF1,
+                WrappingAlgorithm.RSA_OAEP_SHA512_MGF1,
+            )
+        ),
+    )
     _private_wrapping_key = attr.ib(default=None, repr=False, validator=optional(instance_of(RSAPrivateKey)))
     _public_wrapping_key = attr.ib(default=None, repr=False, validator=optional(instance_of(RSAPublicKey)))
 

test/functional/keyrings/raw/test_raw_aes.py

@@ -65,21 +65,14 @@ def sample_encryption_materials():
 
 
 @pytest.mark.parametrize("encryption_materials_samples", sample_encryption_materials())
-@pytest.mark.parametrize("wrapping_algorithm_samples", _WRAPPING_ALGORITHM)
-def test_raw_aes_encryption_decryption(encryption_materials_samples, wrapping_algorithm_samples):
+def test_raw_aes_encryption_decryption(encryption_materials_samples):
 
     # Initializing attributes
     key_namespace = _PROVIDER_ID
     key_name = _KEY_ID
-    _wrapping_algorithm = wrapping_algorithm_samples
 
     # Creating an instance of a raw AES keyring
-    test_raw_aes_keyring = RawAESKeyring(
-        key_namespace=key_namespace,
-        key_name=key_name,
-        wrapping_key=_WRAPPING_KEY,
-        wrapping_algorithm=_wrapping_algorithm,
-    )
+    test_raw_aes_keyring = RawAESKeyring(key_namespace=key_namespace, key_name=key_name, wrapping_key=_WRAPPING_KEY,)
 
     # Call on_encrypt function for the keyring
     encryption_materials = test_raw_aes_keyring.on_encrypt(encryption_materials=encryption_materials_samples)
@@ -101,21 +94,14 @@ def test_raw_aes_encryption_decryption(encryption_materials_samples, wrapping_al
 
 
 @pytest.mark.parametrize("encryption_materials_samples", sample_encryption_materials())
-@pytest.mark.parametrize("wrapping_algorithm_samples", _WRAPPING_ALGORITHM)
-def test_raw_master_key_decrypts_what_raw_keyring_encrypts(encryption_materials_samples, wrapping_algorithm_samples):
+def test_raw_master_key_decrypts_what_raw_keyring_encrypts(encryption_materials_samples):
 
     # Initializing attributes
     key_namespace = _PROVIDER_ID
     key_name = _KEY_ID
-    _wrapping_algorithm = wrapping_algorithm_samples
 
     # Creating an instance of a raw AES keyring
-    test_raw_aes_keyring = RawAESKeyring(
-        key_namespace=key_namespace,
-        key_name=key_name,
-        wrapping_key=_WRAPPING_KEY,
-        wrapping_algorithm=_wrapping_algorithm,
-    )
+    test_raw_aes_keyring = RawAESKeyring(key_namespace=key_namespace, key_name=key_name, wrapping_key=_WRAPPING_KEY,)
 
     # Creating an instance of a raw master key
     test_raw_master_key = RawMasterKey(
@@ -139,21 +125,14 @@ def test_raw_master_key_decrypts_what_raw_keyring_encrypts(encryption_materials_
 
 
 @pytest.mark.parametrize("encryption_materials_samples", sample_encryption_materials())
-@pytest.mark.parametrize("wrapping_algorithm_samples", _WRAPPING_ALGORITHM)
-def test_raw_keyring_decrypts_what_raw_master_key_encrypts(encryption_materials_samples, wrapping_algorithm_samples):
+def test_raw_keyring_decrypts_what_raw_master_key_encrypts(encryption_materials_samples):
 
     # Initializing attributes
     key_namespace = _PROVIDER_ID
     key_name = _KEY_ID
-    _wrapping_algorithm = wrapping_algorithm_samples
 
     # Creating an instance of a raw AES keyring
-    test_raw_aes_keyring = RawAESKeyring(
-        key_namespace=key_namespace,
-        key_name=key_name,
-        wrapping_key=_WRAPPING_KEY,
-        wrapping_algorithm=_wrapping_algorithm,
-    )
+    test_raw_aes_keyring = RawAESKeyring(key_namespace=key_namespace, key_name=key_name, wrapping_key=_WRAPPING_KEY,)
 
     # Creating an instance of a raw master key
     test_raw_master_key = RawMasterKey(

test/functional/keyrings/test_multi.py

@@ -50,12 +50,7 @@
 )
 
 _MULTI_KEYRING_WITH_GENERATOR_AND_CHILDREN = MultiKeyring(
-    generator=RawAESKeyring(
-        key_namespace=_PROVIDER_ID,
-        key_name=_KEY_ID,
-        wrapping_algorithm=WrappingAlgorithm.AES_256_GCM_IV12_TAG16_NO_PADDING,
-        wrapping_key=_WRAPPING_KEY_AES,
-    ),
+    generator=RawAESKeyring(key_namespace=_PROVIDER_ID, key_name=_KEY_ID, wrapping_key=_WRAPPING_KEY_AES,),
     children=[
         RawRSAKeyring(
             key_namespace=_PROVIDER_ID,
@@ -95,12 +90,7 @@
                 public_exponent=65537, key_size=2048, backend=default_backend()
             ),
         ),
-        RawAESKeyring(
-            key_namespace=_PROVIDER_ID,
-            key_name=_KEY_ID,
-            wrapping_algorithm=WrappingAlgorithm.AES_128_GCM_IV12_TAG16_NO_PADDING,
-            wrapping_key=_WRAPPING_KEY_AES,
-        ),
+        RawAESKeyring(key_namespace=_PROVIDER_ID, key_name=_KEY_ID, wrapping_key=_WRAPPING_KEY_AES,),
     ]
 )
 

test/unit/keyrings/raw/test_raw_aes.py

@@ -47,12 +47,7 @@
 
 @pytest.fixture
 def raw_aes_keyring():
-    return RawAESKeyring(
-        key_namespace=_PROVIDER_ID,
-        key_name=_KEY_ID,
-        wrapping_algorithm=WrappingAlgorithm.AES_256_GCM_IV12_TAG16_NO_PADDING,
-        wrapping_key=_WRAPPING_KEY,
-    )
+    return RawAESKeyring(key_namespace=_PROVIDER_ID, key_name=_KEY_ID, wrapping_key=_WRAPPING_KEY,)
 
 
 @pytest.fixture
@@ -113,12 +108,18 @@ def test_valid_parameters(raw_aes_keyring):
 def test_invalid_parameters(key_namespace, key_name, wrapping_algorithm, wrapping_key):
     with pytest.raises(TypeError):
         RawAESKeyring(
-            key_namespace=key_namespace,
-            key_name=key_name,
-            wrapping_algorithm=wrapping_algorithm,
-            wrapping_key=wrapping_key,
+            key_namespace=key_namespace, key_name=key_name, wrapping_key=wrapping_key,
+        )
+
+
+def test_invalid_key_length():
+    with pytest.raises(ValueError) as excinfo:
+        RawAESKeyring(
+            key_namespace=_PROVIDER_ID, key_name=_KEY_ID, wrapping_key=b"012345",
         )
 
+    excinfo.match(r"Invalid wrapping key length. Must be one of \[16, 24, 32\] bytes.")
+
 
 def test_on_encrypt_when_data_encryption_key_given(raw_aes_keyring, patch_generate_data_key):
     test_raw_aes_keyring = raw_aes_keyring

test/unit/keyrings/raw/test_raw_rsa.py

@@ -108,6 +108,24 @@ def test_invalid_parameters(key_namespace, key_name, wrapping_algorithm, private
         )
 
 
+@pytest.mark.parametrize(
+    "wrapping_algorithm",
+    (
+        WrappingAlgorithm.AES_128_GCM_IV12_TAG16_NO_PADDING,
+        WrappingAlgorithm.AES_192_GCM_IV12_TAG16_NO_PADDING,
+        WrappingAlgorithm.AES_256_GCM_IV12_TAG16_NO_PADDING,
+    ),
+)
+def test_invalid_wrapping_algorithm_suite(wrapping_algorithm):
+    with pytest.raises(ValueError):
+        RawRSAKeyring(
+            key_namespace=_PROVIDER_ID,
+            key_name=_KEY_ID,
+            wrapping_algorithm=wrapping_algorithm,
+            private_wrapping_key=raw_rsa_private_key(),
+        )
+
+
 def test_public_and_private_key_not_provided():
     with pytest.raises(TypeError) as exc_info:
         RawRSAKeyring(

test/unit/keyrings/test_multi.py

@@ -96,26 +96,14 @@ def test_parent():
 
 
 def test_keyring_with_generator_but_no_children():
-    generator_keyring = RawAESKeyring(
-        key_namespace=_PROVIDER_ID,
-        key_name=_KEY_ID,
-        wrapping_key=_WRAPPING_KEY_AES,
-        wrapping_algorithm=WrappingAlgorithm.AES_256_GCM_IV12_TAG16_NO_PADDING,
-    )
+    generator_keyring = RawAESKeyring(key_namespace=_PROVIDER_ID, key_name=_KEY_ID, wrapping_key=_WRAPPING_KEY_AES,)
     test_multi_keyring = MultiKeyring(generator=generator_keyring)
     assert test_multi_keyring.generator is generator_keyring
     assert not test_multi_keyring.children
 
 
 def test_keyring_with_children_but_no_generator():
-    children_keyring = [
-        RawAESKeyring(
-            key_namespace=_PROVIDER_ID,
-            key_name=_KEY_ID,
-            wrapping_key=_WRAPPING_KEY_AES,
-            wrapping_algorithm=WrappingAlgorithm.AES_256_GCM_IV12_TAG16_NO_PADDING,
-        )
-    ]
+    children_keyring = [RawAESKeyring(key_namespace=_PROVIDER_ID, key_name=_KEY_ID, wrapping_key=_WRAPPING_KEY_AES,)]
     test_multi_keyring = MultiKeyring(children=children_keyring)
     assert test_multi_keyring.children is children_keyring
     assert test_multi_keyring.generator is None

test/unit/unit_test_utils.py

@@ -256,12 +256,7 @@ def get_decryption_materials_without_data_key():
 
 def get_multi_keyring_with_generator_and_children():
     return MultiKeyring(
-        generator=RawAESKeyring(
-            key_namespace=_PROVIDER_ID,
-            key_name=_KEY_ID,
-            wrapping_algorithm=WrappingAlgorithm.AES_256_GCM_IV12_TAG16_NO_PADDING,
-            wrapping_key=_WRAPPING_KEY_AES,
-        ),
+        generator=RawAESKeyring(key_namespace=_PROVIDER_ID, key_name=_KEY_ID, wrapping_key=_WRAPPING_KEY_AES,),
         children=[
             RawRSAKeyring(
                 key_namespace=_PROVIDER_ID,
@@ -307,12 +302,7 @@ def get_multi_keyring_with_no_generator():
                     public_exponent=65537, key_size=2048, backend=default_backend()
                 ),
             ),
-            RawAESKeyring(
-                key_namespace=_PROVIDER_ID,
-                key_name=_KEY_ID,
-                wrapping_algorithm=WrappingAlgorithm.AES_128_GCM_IV12_TAG16_NO_PADDING,
-                wrapping_key=_WRAPPING_KEY_AES,
-            ),
+            RawAESKeyring(key_namespace=_PROVIDER_ID, key_name=_KEY_ID, wrapping_key=_WRAPPING_KEY_AES,),
         ]
     )
 
@@ -482,10 +472,7 @@ def ephemeral_raw_aes_keyring(wrapping_algorithm=WrappingAlgorithm.AES_256_GCM_I
     if key is None:
         key = os.urandom(key_length)
     return RawAESKeyring(
-        key_namespace="fake",
-        key_name="aes-{}".format(key_length * 8).encode("utf-8"),
-        wrapping_algorithm=wrapping_algorithm,
-        wrapping_key=key,
+        key_namespace="fake", key_name="aes-{}".format(key_length * 8).encode("utf-8"), wrapping_key=key,
     )