diff options
| -rw-r--r-- | cryptography/primitives/block/base.py | 67 |
1 files changed, 29 insertions, 38 deletions
diff --git a/cryptography/primitives/block/base.py b/cryptography/primitives/block/base.py index adade3ca..650e39c1 100644 --- a/cryptography/primitives/block/base.py +++ b/cryptography/primitives/block/base.py @@ -13,13 +13,16 @@ from __future__ import absolute_import, division, print_function -import abc - -import six +from enum import Enum from cryptography.bindings import _default_api +class _Operation(Enum): + encrypt = 0 + decrypt = 1 + + class BlockCipher(object): def __init__(self, cipher, mode, api=None): super(BlockCipher, self).__init__() @@ -32,51 +35,39 @@ class BlockCipher(object): self._api = api def encryptor(self): - return _BlockCipherEncryptionContext(self.cipher, self.mode, self._api) + return _BlockCipherContext(self.cipher, self.mode, self._api, + _Operation.encrypt) def decryptor(self): - return _BlockCipherDecryptionContext(self.cipher, self.mode, self._api) + return _BlockCipherContext(self.cipher, self.mode, self._api, + _Operation.decrypt) -class _BlockCipherContext(six.with_metaclass(abc.ABCMeta)): - def __init__(self, cipher, mode, api): +class _BlockCipherContext(object): + def __init__(self, cipher, mode, api, operation): super(_BlockCipherContext, self).__init__() - self.cipher = cipher - self.mode = mode self._api = api + self._operation = operation + args = (cipher, mode) + if self._operation == _Operation.encrypt: + self._ctx = self._api.create_block_cipher_encrypt_context(*args) + else: + self._ctx = self._api.create_block_cipher_decrypt_context(*args) - def _check_ctx(self): + def update(self, data): if self._ctx is None: raise ValueError("Context was already finalized") - - -class _BlockCipherEncryptionContext(_BlockCipherContext): - def __init__(self, cipher, mode, api): - super(_BlockCipherEncryptionContext, self).__init__(cipher, mode, api) - self._ctx = self._api.create_block_cipher_encrypt_context(cipher, mode) - - def update(self, data): - self._check_ctx() - return self._api.update_encrypt_context(self._ctx, data) + if self._operation == _Operation.encrypt: + return self._api.update_encrypt_context(self._ctx, data) + else: + return self._api.update_decrypt_context(self._ctx, data) def finalize(self): - self._check_ctx() - data = self._api.finalize_encrypt_context(self._ctx) - self._ctx = None - return data - - -class _BlockCipherDecryptionContext(_BlockCipherContext): - def __init__(self, cipher, mode, api): - super(_BlockCipherDecryptionContext, self).__init__(cipher, mode, api) - self._ctx = self._api.create_block_cipher_decrypt_context(cipher, mode) - - def update(self, data): - self._check_ctx() - return self._api.update_decrypt_context(self._ctx, data) - - def finalize(self): - self._check_ctx() - data = self._api.finalize_decrypt_context(self._ctx) + if self._ctx is None: + raise ValueError("Context was already finalized") + if self._operation == _Operation.encrypt: + data = self._api.finalize_encrypt_context(self._ctx) + else: + data = self._api.finalize_decrypt_context(self._ctx) self._ctx = None return data |