implement AES KW with padding (RFC 5649) (#3880)

* implement AES KW with padding (RFC 5649)

fixes #3791

* oops, 2.2

* make sure this is the right valueerror

* more match

* make key padding easier to read

* review feedback

* review feedback

1 files changed, 57 insertions, 0 deletions

diff --git a/src/cryptography/hazmat/primitives/keywrap.py b/src/cryptography/hazmat/primitives/keywrap.py
index 702a6932..3b531318 100644
--- a/src/cryptography/hazmat/primitives/keywrap.py
+++ b/src/cryptography/hazmat/primitives/keywrap.py
@@ -68,6 +68,63 @@ def _unwrap_core(wrapping_key, a, r, backend):
     return a, r
 
 
+def aes_key_wrap_with_padding(wrapping_key, key_to_wrap, backend):
+    if len(wrapping_key) not in [16, 24, 32]:
+        raise ValueError("The wrapping key must be a valid AES key length")
+
+    aiv = b"\xA6\x59\x59\xA6" + struct.pack(">i", len(key_to_wrap))
+    # pad the key to wrap if necessary
+    pad = (8 - (len(key_to_wrap) % 8)) % 8
+    key_to_wrap = key_to_wrap + b"\x00" * pad
+    if len(key_to_wrap) == 8:
+        # RFC 5649 - 4.1 - exactly 8 octets after padding
+        encryptor = Cipher(AES(wrapping_key), ECB(), backend).encryptor()
+        b = encryptor.update(aiv + key_to_wrap)
+        assert encryptor.finalize() == b""
+        return b
+    else:
+        r = [key_to_wrap[i:i + 8] for i in range(0, len(key_to_wrap), 8)]
+        return _wrap_core(wrapping_key, aiv, r, backend)
+
+
+def aes_key_unwrap_with_padding(wrapping_key, wrapped_key, backend):
+    if len(wrapped_key) < 16:
+        raise ValueError("Must be at least 16 bytes")
+
+    if len(wrapping_key) not in [16, 24, 32]:
+        raise ValueError("The wrapping key must be a valid AES key length")
+
+    if len(wrapped_key) == 16:
+        # RFC 5649 - 4.2 - exactly two 64-bit blocks
+        decryptor = Cipher(AES(wrapping_key), ECB(), backend).decryptor()
+        b = decryptor.update(wrapped_key)
+        assert decryptor.finalize() == b""
+        a = b[:8]
+        data = b[8:]
+        n = 1
+    else:
+        r = [wrapped_key[i:i + 8] for i in range(0, len(wrapped_key), 8)]
+        encrypted_aiv = r.pop(0)
+        n = len(r)
+        a, r = _unwrap_core(wrapping_key, encrypted_aiv, r, backend)
+        data = b"".join(r)
+
+    # 1) Check that MSB(32,A) = A65959A6.
+    # 2) Check that 8*(n-1) < LSB(32,A) <= 8*n.  If so, let
+    #    MLI = LSB(32,A).
+    # 3) Let b = (8*n)-MLI, and then check that the rightmost b octets of
+    #    the output data are zero.
+    (mli,) = struct.unpack(">I", a[4:])
+    b = (8 * n) - mli
+    if (
+        not bytes_eq(a[:4], b"\xa6\x59\x59\xa6") or not
+        8 * (n - 1) < mli <= 8 * n or not bytes_eq(data[-b:], b"\x00" * b)
+    ):
+        raise InvalidUnwrap()
+
+    return data[:-b]
+
+
 def aes_key_unwrap(wrapping_key, wrapped_key, backend):
     if len(wrapped_key) < 24:
         raise ValueError("Must be at least 24 bytes")