Merge pull request #112 from reaperhulk/block-cipher-decrypt

Block Cipher Decryption

5 files changed, 110 insertions, 50 deletions

diff --git a/cryptography/bindings/openssl/api.py b/cryptography/bindings/openssl/api.py
index 3d92c144..fedaf9cc 100644
--- a/cryptography/bindings/openssl/api.py
+++ b/cryptography/bindings/openssl/api.py
@@ -142,7 +142,27 @@ class API(object):
                 GetCipherByName("des-ede3-{mode.name}")
             )
 
-    def create_block_cipher_context(self, cipher, mode):
+    def create_block_cipher_encrypt_context(self, cipher, mode):
+        ctx, evp, iv_nonce = self._create_block_cipher_context(cipher, mode)
+        res = self.lib.EVP_EncryptInit_ex(ctx, evp, api.ffi.NULL, cipher.key,
+                                          iv_nonce)
+        assert res != 0
+        # We purposely disable padding here as it's handled higher up in the
+        # API.
+        self.lib.EVP_CIPHER_CTX_set_padding(ctx, 0)
+        return ctx
+
+    def create_block_cipher_decrypt_context(self, cipher, mode):
+        ctx, evp, iv_nonce = self._create_block_cipher_context(cipher, mode)
+        res = self.lib.EVP_DecryptInit_ex(ctx, evp, api.ffi.NULL, cipher.key,
+                                          iv_nonce)
+        assert res != 0
+        # We purposely disable padding here as it's handled higher up in the
+        # API.
+        self.lib.EVP_CIPHER_CTX_set_padding(ctx, 0)
+        return ctx
+
+    def _create_block_cipher_context(self, cipher, mode):
         ctx = self.lib.EVP_CIPHER_CTX_new()
         ctx = self.ffi.gc(ctx, self.lib.EVP_CIPHER_CTX_free)
         evp_cipher = self._cipher_registry[type(cipher), type(mode)](
@@ -156,24 +176,21 @@ class API(object):
         else:
             iv_nonce = self.ffi.NULL
 
-        # TODO: Sometimes this needs to be a DecryptInit, when?
-        res = self.lib.EVP_EncryptInit_ex(
-            ctx, evp_cipher, self.ffi.NULL, cipher.key, iv_nonce
-        )
-        assert res != 0
+        return (ctx, evp_cipher, iv_nonce)
 
-        # We purposely disable padding here as it's handled higher up in the
-        # API.
-        self.lib.EVP_CIPHER_CTX_set_padding(ctx, 0)
-        return ctx
+    def update_encrypt_context(self, ctx, data):
+        block_size = self.lib.EVP_CIPHER_CTX_block_size(ctx)
+        buf = self.ffi.new("unsigned char[]", len(data) + block_size - 1)
+        outlen = self.ffi.new("int *")
+        res = self.lib.EVP_EncryptUpdate(ctx, buf, outlen, data, len(data))
+        assert res != 0
+        return self.ffi.buffer(buf)[:outlen[0]]
 
-    def update_encrypt_context(self, ctx, plaintext):
+    def update_decrypt_context(self, ctx, data):
         block_size = self.lib.EVP_CIPHER_CTX_block_size(ctx)
-        buf = self.ffi.new("unsigned char[]", len(plaintext) + block_size - 1)
+        buf = self.ffi.new("unsigned char[]", len(data) + block_size - 1)
         outlen = self.ffi.new("int *")
-        res = self.lib.EVP_EncryptUpdate(
-            ctx, buf, outlen, plaintext, len(plaintext)
-        )
+        res = self.lib.EVP_DecryptUpdate(ctx, buf, outlen, data, len(data))
         assert res != 0
         return self.ffi.buffer(buf)[:outlen[0]]
 
@@ -187,6 +204,16 @@ class API(object):
         assert res == 1
         return self.ffi.buffer(buf)[:outlen[0]]
 
+    def finalize_decrypt_context(self, ctx):
+        block_size = self.lib.EVP_CIPHER_CTX_block_size(ctx)
+        buf = self.ffi.new("unsigned char[]", block_size)
+        outlen = self.ffi.new("int *")
+        res = self.lib.EVP_DecryptFinal_ex(ctx, buf, outlen)
+        assert res != 0
+        res = self.lib.EVP_CIPHER_CTX_cleanup(ctx)
+        assert res == 1
+        return self.ffi.buffer(buf)[:outlen[0]]
+
     def supports_hash(self, hash_cls):
         return (self.ffi.NULL !=
                 self.lib.EVP_get_digestbyname(hash_cls.name.encode("ascii")))
diff --git a/cryptography/bindings/openssl/evp.py b/cryptography/bindings/openssl/evp.py
index 2bb5b0f7..41df1056 100644
--- a/cryptography/bindings/openssl/evp.py
+++ b/cryptography/bindings/openssl/evp.py
@@ -41,6 +41,11 @@ int EVP_CIPHER_CTX_set_padding(EVP_CIPHER_CTX *, int);
 int EVP_EncryptUpdate(EVP_CIPHER_CTX *, unsigned char *, int *,
                       const unsigned char *, int);
 int EVP_EncryptFinal_ex(EVP_CIPHER_CTX *, unsigned char *, int *);
+int EVP_DecryptInit_ex(EVP_CIPHER_CTX *, const EVP_CIPHER *, ENGINE *,
+                       const unsigned char *, const unsigned char *);
+int EVP_DecryptUpdate(EVP_CIPHER_CTX *, unsigned char *, int *,
+                      const unsigned char *, int);
+int EVP_DecryptFinal_ex(EVP_CIPHER_CTX *, unsigned char *, int *);
 int EVP_CIPHER_CTX_cleanup(EVP_CIPHER_CTX *);
 const EVP_CIPHER *EVP_CIPHER_CTX_cipher(const EVP_CIPHER_CTX *);
 int EVP_CIPHER_block_size(const EVP_CIPHER *);
diff --git a/cryptography/primitives/block/base.py b/cryptography/primitives/block/base.py
index 42c1f799..12b6f626 100644
--- a/cryptography/primitives/block/base.py
+++ b/cryptography/primitives/block/base.py
@@ -13,12 +13,7 @@
 
 from __future__ import absolute_import, division, print_function
 
-from enum import Enum
-
-
-class _Operation(Enum):
-    encrypt = 0
-    decrypt = 1
+from cryptography.primitives import interfaces
 
 
 class BlockCipher(object):
@@ -31,30 +26,49 @@ class BlockCipher(object):
         self.cipher = cipher
         self.mode = mode
         self._api = api
-        self._ctx = api.create_block_cipher_context(cipher, mode)
-        self._operation = None
 
-    def encrypt(self, plaintext):
-        if self._ctx is None:
-            raise ValueError("BlockCipher was already finalized")
+    def encryptor(self):
+        return _CipherEncryptionContext(self.cipher, self.mode, self._api)
+
+    def decryptor(self):
+        return _CipherDecryptionContext(self.cipher, self.mode, self._api)
 
-        if self._operation is None:
-            self._operation = _Operation.encrypt
-        elif self._operation is not _Operation.encrypt:
-            raise ValueError("BlockCipher cannot encrypt when the operation is"
-                             " set to %s" % self._operation.name)
 
-        return self._api.update_encrypt_context(self._ctx, plaintext)
+@interfaces.register(interfaces.CipherContext)
+class _CipherEncryptionContext(object):
+    def __init__(self, cipher, mode, api):
+        super(_CipherEncryptionContext, self).__init__()
+        self._api = api
+        self._ctx = self._api.create_block_cipher_encrypt_context(cipher, mode)
+
+    def update(self, data):
+        if self._ctx is None:
+            raise ValueError("Context was already finalized")
+        return self._api.update_encrypt_context(self._ctx, data)
 
     def finalize(self):
         if self._ctx is None:
-            raise ValueError("BlockCipher was already finalized")
+            raise ValueError("Context was already finalized")
+        data = self._api.finalize_encrypt_context(self._ctx)
+        self._ctx = None
+        return data
+
 
-        if self._operation is _Operation.encrypt:
-            result = self._api.finalize_encrypt_context(self._ctx)
-        else:
-            raise ValueError("BlockCipher cannot finalize the unknown "
-                             "operation %s" % self._operation.name)
+@interfaces.register(interfaces.CipherContext)
+class _CipherDecryptionContext(object):
+    def __init__(self, cipher, mode, api):
+        super(_CipherDecryptionContext, self).__init__()
+        self._api = api
+        self._ctx = self._api.create_block_cipher_decrypt_context(cipher, mode)
 
+    def update(self, data):
+        if self._ctx is None:
+            raise ValueError("Context was already finalized")
+        return self._api.update_decrypt_context(self._ctx, data)
+
+    def finalize(self):
+        if self._ctx is None:
+            raise ValueError("Context was already finalized")
+        data = self._api.finalize_decrypt_context(self._ctx)
         self._ctx = None
-        return result
+        return data
diff --git a/cryptography/primitives/block/modes.py b/cryptography/primitives/block/modes.py
index 43631801..a933c187 100644
--- a/cryptography/primitives/block/modes.py
+++ b/cryptography/primitives/block/modes.py
@@ -16,14 +16,7 @@ from __future__ import absolute_import, division, print_function
 from cryptography.primitives import interfaces
 
 
-def register(iface):
-    def register_decorator(klass):
-        iface.register(klass)
-        return klass
-    return register_decorator
-
-
-@register(interfaces.ModeWithInitializationVector)
+@interfaces.register(interfaces.ModeWithInitializationVector)
 class CBC(object):
     name = "CBC"
 
@@ -36,7 +29,7 @@ class ECB(object):
     name = "ECB"
 
 
-@register(interfaces.ModeWithInitializationVector)
+@interfaces.register(interfaces.ModeWithInitializationVector)
 class OFB(object):
     name = "OFB"
 
@@ -45,7 +38,7 @@ class OFB(object):
         self.initialization_vector = initialization_vector
 
 
-@register(interfaces.ModeWithInitializationVector)
+@interfaces.register(interfaces.ModeWithInitializationVector)
 class CFB(object):
     name = "CFB"
 
@@ -54,7 +47,7 @@ class CFB(object):
         self.initialization_vector = initialization_vector
 
 
-@register(interfaces.ModeWithNonce)
+@interfaces.register(interfaces.ModeWithNonce)
 class CTR(object):
     name = "CTR"
 
diff --git a/cryptography/primitives/interfaces.py b/cryptography/primitives/interfaces.py
index c1fc9910..49c19d0e 100644
--- a/cryptography/primitives/interfaces.py
+++ b/cryptography/primitives/interfaces.py
@@ -18,9 +18,30 @@ import abc
 import six
 
 
+def register(iface):
+    def register_decorator(klass):
+        iface.register(klass)
+        return klass
+    return register_decorator
+
+
 class ModeWithInitializationVector(six.with_metaclass(abc.ABCMeta)):
     pass
 
 
 class ModeWithNonce(six.with_metaclass(abc.ABCMeta)):
     pass
+
+
+class CipherContext(six.with_metaclass(abc.ABCMeta)):
+    @abc.abstractmethod
+    def update(self, data):
+        """
+        update takes bytes and return bytes
+        """
+
+    @abc.abstractmethod
+    def finalize(self):
+        """
+        finalize return bytes
+        """