1 files changed, 65 insertions, 0 deletions
diff --git a/tests/hazmat/primitives/test_keywrap.py b/tests/hazmat/primitives/test_keywrap.py
index f41baedb..8311c2a4 100644
--- a/tests/hazmat/primitives/test_keywrap.py
+++ b/tests/hazmat/primitives/test_keywrap.py
@@ -114,3 +114,68 @@ class TestAESKeyWrap(object):
         # Keys to unwrap must be a multiple of 8 bytes
         with pytest.raises(ValueError):
             keywrap.aes_key_unwrap(b"sixteen_byte_key", b"\x00" * 27, backend)
+
+
+@pytest.mark.supported(
+    only_if=lambda backend: backend.cipher_supported(
+        algorithms.AES(b"\x00" * 16), modes.ECB()
+    ),
+    skip_message="Does not support AES key wrap (RFC 5649) because AES-ECB"
+                 " is unsupported",
+)
+@pytest.mark.requires_backend_interface(interface=CipherBackend)
+class TestAESKeyWrapWithPadding(object):
+    @pytest.mark.parametrize(
+        "params",
+        _load_all_params(
+            os.path.join("keywrap", "kwtestvectors"),
+            ["KWP_AE_128.txt", "KWP_AE_192.txt", "KWP_AE_256.txt"],
+            load_nist_vectors
+        )
+    )
+    def test_wrap(self, backend, params):
+        wrapping_key = binascii.unhexlify(params["k"])
+        key_to_wrap = binascii.unhexlify(params["p"])
+        wrapped_key = keywrap.aes_key_wrap_with_padding(
+            wrapping_key, key_to_wrap, backend
+        )
+        assert params["c"] == binascii.hexlify(wrapped_key)
+
+    @pytest.mark.parametrize(
+        "params",
+        _load_all_params(
+            os.path.join("keywrap", "kwtestvectors"),
+            ["KWP_AD_128.txt", "KWP_AD_192.txt", "KWP_AD_256.txt"],
+            load_nist_vectors
+        )
+    )
+    def test_unwrap(self, backend, params):
+        wrapping_key = binascii.unhexlify(params["k"])
+        wrapped_key = binascii.unhexlify(params["c"])
+        if params.get("fail") is True:
+            with pytest.raises(keywrap.InvalidUnwrap):
+                keywrap.aes_key_unwrap_with_padding(
+                    wrapping_key, wrapped_key, backend
+                )
+        else:
+            unwrapped_key = keywrap.aes_key_unwrap_with_padding(
+                wrapping_key, wrapped_key, backend
+            )
+            assert params["p"] == binascii.hexlify(unwrapped_key)
+
+    def test_unwrap_invalid_wrapped_key_length(self, backend):
+        # Keys to unwrap must be at least 16 bytes
+        with pytest.raises(ValueError, match='Must be at least 16 bytes'):
+            keywrap.aes_key_unwrap_with_padding(
+                b"sixteen_byte_key", b"\x00" * 15, backend
+            )
+
+    def test_wrap_invalid_key_length(self, backend):
+        with pytest.raises(ValueError, match='must be a valid AES key length'):
+            keywrap.aes_key_wrap_with_padding(b"badkey", b"\x00", backend)
+
+    def test_unwrap_invalid_key_length(self, backend):
+        with pytest.raises(ValueError, match='must be a valid AES key length'):
+            keywrap.aes_key_unwrap_with_padding(
+                b"badkey", b"\x00" * 16, backend
+            )