add AES128/AES256 classes (#7542)

These let developers be more explicit about the allowable key lengths
for an AES key and make auditing the codebase a bit easier.

But that's not really why we're adding them. In some upcoming
serialization features we need to be able to specify AES 128 vs AES 256
and the current class doesn't work for that since it computes key
length from the key you provide it when instantiating the class.
That's incompatible with serialization where the key is derived
later in the process. C'est la vie.

CHANGELOG.rst

@@ -46,14 +46,21 @@ Changelog
   :class:`~cryptography.hazmat.primitives.kdf.kbkdf.KBKDFCMAC` now support
   :attr:`~cryptography.hazmat.primitives.kdf.kbkdf.CounterLocation.MiddleFixed`
   counter location.
-* Fixed :rfc:`4514` name parsing to reverse the order of the RDNs according 
-  to the section 2.1 of the RFC, affecting method 
+* Fixed :rfc:`4514` name parsing to reverse the order of the RDNs according
+  to the section 2.1 of the RFC, affecting method
   :meth:`~cryptography.x509.Name.from_rfc4514_string`.
 * It is now possible to customize some aspects of encryption when serializing
   private keys, using
   :meth:`~cryptography.hazmat.primitives.serialization.PrivateFormat.encryption_builder`.
 * Removed several legacy symbols from our OpenSSL bindings. Users of pyOpenSSL
   versions older than 22.0 will need to upgrade.
+* Added
+  :class:`~cryptography.hazmat.primitives.ciphers.algorithms.AES128` and
+  :class:`~cryptography.hazmat.primitives.ciphers.algorithms.AES256` classes.
+  These classes do not replace
+  :class:`~cryptography.hazmat.primitives.ciphers.algorithms.AES` (which
+  allows all AES key lengths), but are intended for applications where
+  developers want to be explicit about key length.
 
 .. _v37-0-4:
 

docs/hazmat/primitives/symmetric-encryption.rst

@@ -95,6 +95,28 @@ Algorithms
         ``192``, or ``256`` :term:`bits` long.
     :type key: :term:`bytes-like`
 
+.. class:: AES128(key)
+
+    .. versionadded:: 38.0.0
+
+    An AES class that only accepts 128 bit keys. This is identical to the
+    standard ``AES`` class except that it will only accept a single key length.
+
+    :param key: The secret key. This must be kept secret. ``128``
+        :term:`bits` long.
+    :type key: :term:`bytes-like`
+
+.. class:: AES256(key)
+
+    .. versionadded:: 38.0.0
+
+    An AES class that only accepts 256 bit keys. This is identical to the
+    standard ``AES`` class except that it will only accept a single key length.
+
+    :param key: The secret key. This must be kept secret. ``256``
+        :term:`bits` long.
+    :type key: :term:`bytes-like`
+
 .. class:: Camellia(key)
 
     Camellia is a block cipher approved for use by `CRYPTREC`_ and ISO/IEC.

src/cryptography/hazmat/backends/openssl/backend.py

@@ -90,6 +90,8 @@ from cryptography.hazmat.primitives.ciphers import (
 )
 from cryptography.hazmat.primitives.ciphers.algorithms import (
     AES,
+    AES128,
+    AES256,
     ARC4,
     Camellia,
     ChaCha20,
@@ -378,12 +380,15 @@ class Backend:
         self._cipher_registry[cipher_cls, mode_cls] = adapter
 
     def _register_default_ciphers(self) -> None:
-        for mode_cls in [CBC, CTR, ECB, OFB, CFB, CFB8, GCM]:
-            self.register_cipher_adapter(
-                AES,
-                mode_cls,
-                GetCipherByName("{cipher.name}-{cipher.key_size}-{mode.name}"),
-            )
+        for cipher_cls in [AES, AES128, AES256]:
+            for mode_cls in [CBC, CTR, ECB, OFB, CFB, CFB8, GCM]:
+                self.register_cipher_adapter(
+                    cipher_cls,
+                    mode_cls,
+                    GetCipherByName(
+                        "{cipher.name}-{cipher.key_size}-{mode.name}"
+                    ),
+                )
         for mode_cls in [CBC, CTR, ECB, OFB, CFB]:
             self.register_cipher_adapter(
                 Camellia,

src/cryptography/hazmat/primitives/ciphers/algorithms.py

@@ -38,6 +38,26 @@ class AES(CipherAlgorithm, BlockCipherAlgorithm):
         return len(self.key) * 8
 
 
+class AES128(CipherAlgorithm, BlockCipherAlgorithm):
+    name = "AES"
+    block_size = 128
+    key_sizes = frozenset([128])
+    key_size = 128
+
+    def __init__(self, key: bytes):
+        self.key = _verify_key_size(self, key)
+
+
+class AES256(CipherAlgorithm, BlockCipherAlgorithm):
+    name = "AES"
+    block_size = 128
+    key_sizes = frozenset([256])
+    key_size = 256
+
+    def __init__(self, key: bytes):
+        self.key = _verify_key_size(self, key)
+
+
 class Camellia(CipherAlgorithm, BlockCipherAlgorithm):
     name = "camellia"
     block_size = 128

src/cryptography/hazmat/primitives/ciphers/modes.py

@@ -12,6 +12,7 @@ from cryptography.hazmat.primitives._cipheralgorithm import (
     BlockCipherAlgorithm,
     CipherAlgorithm,
 )
+from cryptography.hazmat.primitives.ciphers import algorithms
 
 
 class Mode(metaclass=abc.ABCMeta):
@@ -135,6 +136,12 @@ class XTS(ModeWithTweak):
         return self._tweak
 
     def validate_for_algorithm(self, algorithm: CipherAlgorithm) -> None:
+        if isinstance(algorithm, (algorithms.AES128, algorithms.AES256)):
+            raise TypeError(
+                "The AES128 and AES256 classes do not support XTS, please use "
+                "the standard AES class instead."
+            )
+
         if algorithm.key_size not in (256, 512):
             raise ValueError(
                 "The XTS specification requires a 256-bit key for AES-128-XTS"

tests/hazmat/primitives/test_aes.py

@@ -73,6 +73,13 @@ class TestAESModeXTS:
         with pytest.raises(ValueError, match="duplicated keys"):
             cipher.encryptor()
 
+    def test_xts_unsupported_with_aes128_aes256_classes(self):
+        with pytest.raises(TypeError):
+            base.Cipher(algorithms.AES128(b"0" * 16), modes.XTS(b"\x00" * 16))
+
+        with pytest.raises(TypeError):
+            base.Cipher(algorithms.AES256(b"0" * 32), modes.XTS(b"\x00" * 16))
+
 
 @pytest.mark.supported(
     only_if=lambda backend: backend.cipher_supported(
@@ -274,3 +281,28 @@ def test_buffer_protocol_alternate_modes(mode, backend):
     dec = cipher.decryptor()
     pt = dec.update(ct) + dec.finalize()
     assert pt == data
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [
+        modes.ECB(),
+        modes.CBC(bytearray(b"\x00" * 16)),
+        modes.CTR(bytearray(b"\x00" * 16)),
+        modes.OFB(bytearray(b"\x00" * 16)),
+        modes.CFB(bytearray(b"\x00" * 16)),
+        modes.CFB8(bytearray(b"\x00" * 16)),
+    ],
+)
+@pytest.mark.parametrize("alg_cls", [algorithms.AES128, algorithms.AES256])
+def test_alternate_aes_classes(mode, alg_cls, backend):
+    alg = alg_cls(b"0" * (alg_cls.key_size // 8))
+    if not backend.cipher_supported(alg, mode):
+        pytest.skip("AES in {} mode not supported".format(mode.name))
+    data = bytearray(b"sixteen_byte_msg")
+    cipher = base.Cipher(alg, mode, backend)
+    enc = cipher.encryptor()
+    ct = enc.update(data) + enc.finalize()
+    dec = cipher.decryptor()
+    pt = dec.update(ct) + dec.finalize()
+    assert pt == data

tests/hazmat/primitives/test_aes_gcm.py

@@ -225,3 +225,15 @@ class TestAESModeGCM:
 
         decryptor.finalize_with_tag(tag)
         assert pt == payload
+
+    @pytest.mark.parametrize("alg", [algorithms.AES128, algorithms.AES256])
+    def test_alternate_aes_classes(self, alg, backend):
+        data = bytearray(b"sixteen_byte_msg")
+        cipher = base.Cipher(
+            alg(b"0" * (alg.key_size // 8)), modes.GCM(b"\x00" * 12), backend
+        )
+        enc = cipher.encryptor()
+        ct = enc.update(data) + enc.finalize()
+        dec = cipher.decryptor()
+        pt = dec.update(ct) + dec.finalize_with_tag(enc.tag)
+        assert pt == data