From ad89af11b234524e543284d3095a15905518ad70 Mon Sep 17 00:00:00 2001
From: Donald Stufft <donald@stufft.io>
Date: Sat, 10 Aug 2013 16:01:56 -0400
Subject: Ensure that AES gets a proper key size

---
 cryptography/primitives/block/ciphers.py | 6 ++++++
 tests/primitives/test_ciphers.py         | 4 ++++
 2 files changed, 10 insertions(+)

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))
-- 
cgit v1.2.3