From 48ddb5373a0a8b0481ed1c72ece36b5973fd3da5 Mon Sep 17 00:00:00 2001
From: Victor Wagner <vitus@wagner.pp.ru>
Date: Sun, 8 Jun 2014 14:45:37 +0400
Subject: [PATCH] Check key and iv length. Handle variable key length ciphers

---
 ctypescrypto/cipher.py | 17 +++++++++++++++++
 tests/testcipher.py    | 22 +++++++++++++++++++++-
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/ctypescrypto/cipher.py b/ctypescrypto/cipher.py
index 1ea4f79..91c5283 100644
--- a/ctypescrypto/cipher.py
+++ b/ctypescrypto/cipher.py
@@ -101,6 +101,10 @@ class Cipher:
 
 		"""
 		self._clean_ctx()
+		# Check key and iv length
+		if key is None:
+			raise ValueError("No key specified")
+
 		key_ptr = c_char_p(key)
 		iv_ptr = c_char_p(iv)
 		self.ctx = libcrypto.EVP_CIPHER_CTX_new()
@@ -111,6 +115,19 @@ class Cipher:
 			enc = 1
 		else: 
 			enc = 0
+		if not iv is None and len(iv) != cipher_type.iv_length():
+			raise ValueError("Invalid IV length for this algorithm")
+			
+		if len(key) != cipher_type.key_length():
+			if (cipher_type.flags() & 8) != 0:
+				# Variable key length cipher.
+				result = libcrypto.EVP_CipherInit_ex(self.ctx,cipher_type.cipher,None,None,None,c_int(enc))
+				result=libcrypto.EVP_CIPHER_CTX_set_key_length(self.ctx,len(key))
+				if result == 0:
+					self._clean_ctx()
+					raise CipherError("Unable to set key length")
+			else:
+				raise ValueError("Invalid key length for this algorithm")
 		result = libcrypto.EVP_CipherInit_ex(self.ctx, cipher_type.cipher, None, key_ptr, iv_ptr, c_int(enc))
 		if result == 0:
 			self._clean_ctx()
diff --git a/tests/testcipher.py b/tests/testcipher.py
index ae74d41..de426cb 100644
--- a/tests/testcipher.py
+++ b/tests/testcipher.py
@@ -64,8 +64,28 @@ class TestEncryptDecrypt(unittest.TestCase):
 		enc=c.update(data)+c.finish()
 		# See if padding is added by default
 		self.assertEqual(len(enc),len(data))
-		d=cipher.new("AES-256-OFB",decryptkey)
+		d=cipher.new("AES-256-OFB",decryptkey,encrypt=False)
 		dec=d.update(enc)+d.finish()
 		self.assertEqual(data,dec)
+	def test_wrong_keylength(self):
+		data="sdfsdfxxx"
+		key="abcdabcd"
+		with self.assertRaises(ValueError):
+			c=cipher.new("AES-128-OFB",key)
+	def test_wrong_ivlength(self):
+		key="abcdabcdabcdabcd"
+		iv="xxxxx"
+		with self.assertRaises(ValueError):
+			c=cipher.new("AES-128-OFB",key,iv=iv)
+	def test_variable_keylength(self):
+		encryptkey="abcdefabcdef"
+		data="asdfsdfsdfsdfsdfsdfsdfsdf"
+		iv="abcdefgh"
+		c=cipher.new("bf-ofb",encryptkey,iv=iv)
+		ciphertext=c.update(data)+c.finish()
+		decryptkey=encryptkey[0:5]+encryptkey[5:]
+		d=cipher.new("bf-ofb",decryptkey,iv=iv)
+		deciph=d.update(ciphertext)+d.finish()
+		self.assertEqual(deciph,data)
 if __name__ == '__main__':
 	unittest.main()
-- 
2.39.5