1 files changed, 39 insertions, 27 deletions
diff --git a/cryptography/primitives/block/base.py b/cryptography/primitives/block/base.py
index 50e9e9e5..12b6f626 100644
--- a/cryptography/primitives/block/base.py
+++ b/cryptography/primitives/block/base.py
@@ -13,14 +13,7 @@
 
 from __future__ import absolute_import, division, print_function
 
-from enum import Enum
-
-from cryptography.bindings import _default_api
-
-
-class _Operation(Enum):
-    encrypt = 0
-    decrypt = 1
+from cryptography.primitives import interfaces
 
 
 class BlockCipher(object):
@@ -28,35 +21,54 @@ class BlockCipher(object):
         super(BlockCipher, self).__init__()
 
         if api is None:
-            api = _default_api
+            from cryptography.bindings import _default_api as api
 
         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)
 
-        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)
+    def decryptor(self):
+        return _CipherDecryptionContext(self.cipher, self.mode, self._api)
 
-        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