_AEADCipherContext refactor
* No longer extends _CipherContext
* Remove _tag from _CipherContext
* This change duplicates a small amount of code from _CipherContext
diff --git a/cryptography/hazmat/primitives/ciphers/base.py b/cryptography/hazmat/primitives/ciphers/base.py
index 3a27030..89e5654 100644
--- a/cryptography/hazmat/primitives/ciphers/base.py
+++ b/cryptography/hazmat/primitives/ciphers/base.py
@@ -60,6 +60,25 @@
class _CipherContext(object):
def __init__(self, ctx):
self._ctx = ctx
+
+ 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._ctx = None
+ return data
+
+
+@utils.register_interface(interfaces.AEADCipherContext)
+@utils.register_interface(interfaces.CipherContext)
+class _AEADCipherContext(object):
+ def __init__(self, ctx):
+ self._ctx = ctx
self._tag = None
def update(self, data):
@@ -75,10 +94,6 @@
self._ctx = None
return data
-
-@utils.register_interface(interfaces.AEADCipherContext)
-@utils.register_interface(interfaces.CipherContext)
-class _AEADCipherContext(_CipherContext):
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 9845555..2a99cab 100644
--- a/tests/hazmat/primitives/utils.py
+++ b/tests/hazmat/primitives/utils.py
@@ -338,3 +338,7 @@
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()