diff options
| -rw-r--r-- | cryptography/primitives/block/ciphers.py | 6 | ||||
| -rw-r--r-- | tests/primitives/test_ciphers.py | 4 |
2 files changed, 10 insertions, 0 deletions
diff --git a/cryptography/primitives/block/ciphers.py b/cryptography/primitives/block/ciphers.py index f40fc2a1..2bf936ee 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 31b4275e..891c5cf8 100644 --- a/tests/primitives/test_ciphers.py +++ b/tests/primitives/test_ciphers.py @@ -27,3 +27,7 @@ class TestAES(object): 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)) |