diff options
| -rw-r--r-- | cryptography/hazmat/primitives/ciphers/base.py | 21 | ||||
| -rw-r--r-- | tests/hazmat/primitives/utils.py | 4 |
2 files changed, 22 insertions, 3 deletions
diff --git a/cryptography/hazmat/primitives/ciphers/base.py b/cryptography/hazmat/primitives/ciphers/base.py index 3a27030a..89e56547 100644 --- a/cryptography/hazmat/primitives/ciphers/base.py +++ b/cryptography/hazmat/primitives/ciphers/base.py @@ -60,7 +60,6 @@ class Cipher(object): class _CipherContext(object): def __init__(self, ctx): self._ctx = ctx - self._tag = None def update(self, data): if self._ctx is None: @@ -71,14 +70,30 @@ class _CipherContext(object): if self._ctx is None: raise AlreadyFinalized("Context was already finalized") data = self._ctx.finalize() - self._tag = self._ctx._tag self._ctx = None return data @utils.register_interface(interfaces.AEADCipherContext) @utils.register_interface(interfaces.CipherContext) -class _AEADCipherContext(_CipherContext): +class _AEADCipherContext(object): + def __init__(self, ctx): + self._ctx = ctx + self._tag = None + + def update(self, data): + if self._ctx is None: + raise AlreadyFinalized("Context was already finalized") + return self._ctx.update(data) + + def finalize(self): + if self._ctx is None: + raise AlreadyFinalized("Context was already finalized") + data = self._ctx.finalize() + self._tag = self._ctx._tag + self._ctx = None + return data + def add_data(self, data): if self._ctx is None: raise AlreadyFinalized("Context was already finalized") diff --git a/tests/hazmat/primitives/utils.py b/tests/hazmat/primitives/utils.py index 98455556..2a99cab9 100644 --- a/tests/hazmat/primitives/utils.py +++ b/tests/hazmat/primitives/utils.py @@ -338,3 +338,7 @@ def aead_use_after_finalize_test(backend, cipher_factory, mode_factory, encryptor.finalize() with pytest.raises(AlreadyFinalized): encryptor.add_data(b"b" * 16) + with pytest.raises(AlreadyFinalized): + encryptor.update(b"b" * 16) + with pytest.raises(AlreadyFinalized): + encryptor.finalize() |