1 files changed, 34 insertions, 44 deletions

diff --git a/src/mailman_pgp/pgp/wrapper.py b/src/mailman_pgp/pgp/wrapper.py
index b8ed264..f746229 100644
--- a/src/mailman_pgp/pgp/wrapper.py
+++ b/src/mailman_pgp/pgp/wrapper.py
@@ -15,11 +15,11 @@
 # this program.  If not, see <http://www.gnu.org/licenses/>.
 
 """A combined PGP/MIME + inline PGP wrapper."""
-import copy
 
 from pgpy.errors import PGPError
 from public import public
 
+from mailman_pgp.pgp.base import BaseWrapper
 from mailman_pgp.pgp.inline import InlineWrapper
 from mailman_pgp.pgp.mime import MIMEWrapper
 from mailman_pgp.pgp.mime_multisig import MIMEMultiSigWrapper
@@ -27,22 +27,25 @@ from mailman_pgp.utils.pgp import verifies
 
 
 @public
-class PGPWrapper():
+class PGPWrapper(BaseWrapper):
     """A combined PGP/MIME + inline PGP wrapper."""
 
-    def __init__(self, msg, default=MIMEWrapper):
+    def __init__(self, msg, copy=False, default=MIMEWrapper):
         """
         Wrap the given message.
 
         :param msg: The message to wrap.
         :type msg: mailman.email.message.Message
-        :param default:
+        :param copy: Whether to copy the message when wrapping.
+        :type copy: bool
+        :param default: The wrapper class used for active operations (sign,
+                        encrypt, attach_keys, attach_revocs)
         :type default: Type[MIMEWrapper|MIMEMultiSigWrapper|InlineWrapper]
         """
-        self.msg = msg
-        self.mime = MIMEWrapper(msg)
-        self.inline = InlineWrapper(msg)
-        self.multisig = MIMEMultiSigWrapper(msg)
+        super().__init__(msg, copy)
+        self.mime = MIMEWrapper(self.msg)
+        self.inline = InlineWrapper(self.msg)
+        self.multisig = MIMEMultiSigWrapper(self.msg)
         self.wrappers = (self.mime, self.inline, self.multisig)
         if default is MIMEWrapper:
             self.default = self.mime
@@ -56,6 +59,10 @@ class PGPWrapper():
                              MIMEMultiSigWrapper.__name__ + ' ' +
                              InlineWrapper.__name__ + '.')
 
+    def _rewrap(self, wrapper):
+        if wrapper is not None:
+            return PGPWrapper(wrapper.msg, default=self.default.__class__)
+
     def get_payload(self):
         return self.default.get_payload()
 
@@ -111,10 +118,10 @@ class PGPWrapper():
 
         :param key: The key to sign with.
         :type key: pgpy.PGPKey
-        :return: The signed message.
-        :rtype: mailman.email.message.Message
+        :return:
+        :rtype: PGPWrapper
         """
-        return self.default.sign(key, **kwargs)
+        return self._rewrap(self.default.sign(key, **kwargs))
 
     def verify(self, key):
         """
@@ -172,12 +179,10 @@ class PGPWrapper():
 
         :param keys: The key/s to encrypt with.
         :type keys: pgpy.PGPKey
-        :param cipher: The symmetric cipher to use.
-        :type cipher: SymmetricKeyAlgorithm
-        :return: The encrypted message.
-        :rtype: mailman.email.message.Message
+        :return:
+        :rtype: PGPWrapper
         """
-        return self.default.encrypt(*keys, **kwargs)
+        return self._rewrap(self.default.encrypt(*keys, **kwargs))
 
     def decrypt(self, key):
         """
@@ -185,16 +190,18 @@ class PGPWrapper():
 
         :param key: The key to decrypt with.
         :type key: pgpy.PGPKey
-        :return: The decrypted message.
         :raises: pgpy.errors.PGPError
-        :rtype: mailman.email.message.Message
+        :return:
+        :rtype: PGPWrapper
         """
+        result = None
         if self.mime.is_encrypted():
-            return self.mime.decrypt(key)
+            result = self.mime.decrypt(key)
         elif self.multisig.is_encrypted():
-            return self.multisig.decrypt(key)
+            result = self.multisig.decrypt(key)
         elif self.inline.is_encrypted():
-            return self.inline.decrypt(key)
+            result = self.inline.decrypt(key)
+        return self._rewrap(result)
 
     def try_decrypt(self, key):
         """
@@ -204,12 +211,12 @@ class PGPWrapper():
         :type key: pgpy.PGPKey
         :return: The decrypted message, if successfully decrypted,
                  else original message.
-        :rtype: mailman.email.message.Message
+        :rtype: PGPWrapper
         """
         try:
-            return self.decrypt(key)
+            return self._rewrap(self.decrypt(key))
         except PGPError:
-            return copy.deepcopy(self.msg)
+            return self
 
     def sign_encrypt(self, key, *keys, **kwargs):
         """
@@ -223,27 +230,10 @@ class PGPWrapper():
         :type hash: pgpy.constants.HashAlgorithm
         :param cipher:
         :type cipher: pgpy.constants.SymmetricKeyAlgorithm
-        :return: The signed + encrypted message.
-        :rtype: mailman.email.message.Message
-        """
-        return self.default.sign_encrypt(key, *keys, **kwargs)
-
-    def sign_then_encrypt(self, key, *keys, **kwargs):
-        """
-        Sign then encrypt the message.
-
-        :param key: The key to sign with.
-        :type key: pgpy.PGPKey
-        :param keys: The key/s to encrypt with.
-        :type keys: pgpy.PGPKey
-        :param hash:
-        :type hash: pgpy.constants.HashAlgorithm
-        :param cipher:
-        :type cipher: pgpy.constants.SymmetricKeyAlgorithm
-        :return: The signed + encrypted message.
-        :rtype: mailman.email.message.Message
+        :return:
+        :rtype: PGPWrapper
         """
-        return self.default.sign_then_encrypt(key, *keys, **kwargs)
+        return self._rewrap(self.default.sign_encrypt(key, *keys, **kwargs))
 
     def is_keys(self):
         """