Use an enum for determining BlockCipher operation

4 files changed, 25 insertions, 9 deletions

diff --git a/cryptography/primitives/block/base.py b/cryptography/primitives/block/base.py
index 207c83d9..2018138c 100644
--- a/cryptography/primitives/block/base.py
+++ b/cryptography/primitives/block/base.py
@@ -12,9 +12,16 @@
 # limitations under the License.
 
 # TODO: which binding is used should be an option somewhere
+from enum import Enum
+
 from cryptography.bindings.openssl import api
 
 
+class _Operation(Enum):
+    encrypt = "encrypt"
+    decrypt = "decrypt"
+
+
 class BlockCipher(object):
     def __init__(self, cipher, mode):
         super(BlockCipher, self).__init__()
@@ -34,10 +41,10 @@ class BlockCipher(object):
             raise ValueError("BlockCipher was already finalized")
 
         if self._operation is None:
-            self._operation = "encrypt"
-        elif self._operation != "encrypt":
+            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)
+                             " set to %s" % self._operation.name)
 
         return api.update_encrypt_context(self._ctx, plaintext)
 
@@ -45,11 +52,11 @@ class BlockCipher(object):
         if self._ctx is None:
             raise ValueError("BlockCipher was already finalized")
 
-        if self._operation == "encrypt":
+        if self._operation is _Operation.encrypt:
             result = api.finalize_encrypt_context(self._ctx)
         else:
             raise ValueError("BlockCipher cannot finalize the unknown "
-                             "operation %s" % self._operation)
+                             "operation %s" % self._operation.name)
 
         self._ctx = None
         return result
diff --git a/setup.py b/setup.py
index c3b754f3..8fd8a724 100644
--- a/setup.py
+++ b/setup.py
@@ -10,12 +10,19 @@
 # implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import sys
 
 from setuptools import setup
 
+install_requires = [
+    "cffi>=0.6",
+]
+
+if sys.version_info[:2] < (3, 4):
+    install_requires += ["enum34"]
 
 setup(
     name="cryptography",
     license="Apache License, Version 2.0",
-    install_requires=["cffi>=0.6"],
+    install_requires=install_requires,
 )
diff --git a/tests/primitives/test_block.py b/tests/primitives/test_block.py
index 799e5f03..9059a886 100644
--- a/tests/primitives/test_block.py
+++ b/tests/primitives/test_block.py
@@ -13,9 +13,11 @@
 
 import binascii
 
+import pretend
 import pytest
 
 from cryptography.primitives.block import BlockCipher, ciphers, modes
+from cryptography.primitives.block.base import _Operation
 
 
 class TestBlockCipher(object):
@@ -43,7 +45,7 @@ class TestBlockCipher(object):
             ciphers.AES(binascii.unhexlify(b"0" * 32)),
             modes.CBC(binascii.unhexlify(b"0" * 32))
         )
-        cipher._operation = "decrypt"
+        cipher._operation = _Operation.decrypt
 
         with pytest.raises(ValueError):
             cipher.encrypt(b"b" * 16)
@@ -53,7 +55,7 @@ class TestBlockCipher(object):
             ciphers.AES(binascii.unhexlify(b"0" * 32)),
             modes.CBC(binascii.unhexlify(b"0" * 32))
         )
-        cipher._operation = "wat"
+        cipher._operation = pretend.stub(name="wat")
 
         with pytest.raises(ValueError):
             cipher.encrypt(b"b" * 16)
diff --git a/tox.ini b/tox.ini
index 81c66758..80e0a15c 100644
--- a/tox.ini
+++ b/tox.ini
@@ -2,7 +2,7 @@
 envlist = py26,py27,pypy,py32,py33,docs,pep8
 
 [testenv]
-deps = pytest-cov
+deps = pytest-cov pretend
 commands = py.test --cov=cryptography/ --cov=tests/
 
 [testenv:docs]