Ensure that AES gets a proper key size
diff --git a/cryptography/primitives/block/ciphers.py b/cryptography/primitives/block/ciphers.py
index f40fc2a..2bf936e 100644
--- a/cryptography/primitives/block/ciphers.py
+++ b/cryptography/primitives/block/ciphers.py
@@ -14,11 +14,17 @@
class AES(object):
name = "AES"
+ key_sizes = {128, 192, 256}
def __init__(self, key):
super(AES, self).__init__()
self.key = key
+ if self.key_size not in self.key_sizes:
+ raise ValueError("Invalid key size (%s) for %s".format(
+ self.key_size, self.name
+ ))
+
@property
def key_size(self):
return len(self.key) * 8
diff --git a/tests/primitives/test_ciphers.py b/tests/primitives/test_ciphers.py
index 31b4275..891c5cf 100644
--- a/tests/primitives/test_ciphers.py
+++ b/tests/primitives/test_ciphers.py
@@ -27,3 +27,7 @@
def test_key_size(self, key, keysize):
cipher = AES(binascii.unhexlify(key))
assert cipher.key_size == keysize
+
+ def test_invalid_key_size(self):
+ with pytest.raises(ValueError):
+ AES(binascii.unhexlify(b"0" * 12))