Merge pull request #286 from alex/code-duplication

Rreduce code duplication

1 files changed, 16 insertions, 43 deletions

diff --git a/cryptography/hazmat/primitives/ciphers/algorithms.py b/cryptography/hazmat/primitives/ciphers/algorithms.py
index 75a87265..a206b273 100644
--- a/cryptography/hazmat/primitives/ciphers/algorithms.py
+++ b/cryptography/hazmat/primitives/ciphers/algorithms.py
@@ -17,6 +17,15 @@ from cryptography import utils
 from cryptography.hazmat.primitives import interfaces
 
 
+def _verify_key_size(algorithm, key):
+    # Verify that the key size matches the expected key size
+    if len(key) * 8 not in algorithm.key_sizes:
+        raise ValueError("Invalid key size ({0}) for {1}".format(
+            len(key) * 8, algorithm.name
+        ))
+    return key
+
+
 @utils.register_interface(interfaces.CipherAlgorithm)
 class AES(object):
     name = "AES"
@@ -24,13 +33,7 @@ class AES(object):
     key_sizes = frozenset([128, 192, 256])
 
     def __init__(self, key):
-        self.key = key
-
-        # Verify that the key size matches the expected key size
-        if self.key_size not in self.key_sizes:
-            raise ValueError("Invalid key size ({0}) for {1}".format(
-                self.key_size, self.name
-            ))
+        self.key = _verify_key_size(self, key)
 
     @property
     def key_size(self):
@@ -44,13 +47,7 @@ class Camellia(object):
     key_sizes = frozenset([128, 192, 256])
 
     def __init__(self, key):
-        self.key = key
-
-        # Verify that the key size matches the expected key size
-        if self.key_size not in self.key_sizes:
-            raise ValueError("Invalid key size ({0}) for {1}".format(
-                self.key_size, self.name
-            ))
+        self.key = _verify_key_size(self, key)
 
     @property
     def key_size(self):
@@ -68,13 +65,7 @@ class TripleDES(object):
             key += key + key
         elif len(key) == 16:
             key += key[:8]
-        self.key = key
-
-        # Verify that the key size matches the expected key size
-        if self.key_size not in self.key_sizes:
-            raise ValueError("Invalid key size ({0}) for {1}".format(
-                self.key_size, self.name
-            ))
+        self.key = _verify_key_size(self, key)
 
     @property
     def key_size(self):
@@ -88,13 +79,7 @@ class Blowfish(object):
     key_sizes = frozenset(range(32, 449, 8))
 
     def __init__(self, key):
-        self.key = key
-
-        # Verify that the key size matches the expected key size
-        if self.key_size not in self.key_sizes:
-            raise ValueError("Invalid key size ({0}) for {1}".format(
-                self.key_size, self.name
-            ))
+        self.key = _verify_key_size(self, key)
 
     @property
     def key_size(self):
@@ -105,16 +90,10 @@ class Blowfish(object):
 class CAST5(object):
     name = "CAST5"
     block_size = 64
-    key_sizes = frozenset([40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128])
+    key_sizes = frozenset(range(40, 129, 8))
 
     def __init__(self, key):
-        self.key = key
-
-        # Verify that the key size matches the expected key size
-        if self.key_size not in self.key_sizes:
-            raise ValueError("Invalid key size ({0}) for {1}".format(
-                self.key_size, self.name
-            ))
+        self.key = _verify_key_size(self, key)
 
     @property
     def key_size(self):
@@ -128,13 +107,7 @@ class ARC4(object):
     key_sizes = frozenset([40, 56, 64, 80, 128, 192, 256])
 
     def __init__(self, key):
-        self.key = key
-
-        # Verify that the key size matches the expected key size
-        if self.key_size not in self.key_sizes:
-            raise ValueError("Invalid key size ({0}) for {1}".format(
-                self.key_size, self.name
-            ))
+        self.key = _verify_key_size(self, key)
 
     @property
     def key_size(self):