_AEADCipherContext refactor

* No longer extends _CipherContext
* Remove _tag from _CipherContext
* This change duplicates a small amount of code from _CipherContext

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()