feature/wrappers-modify

1 files changed, 34 insertions, 43 deletions

diff --git a/src/mailman_pgp/pgp/inline.py b/src/mailman_pgp/pgp/inline.py
index fa8f878..410078c 100644
--- a/src/mailman_pgp/pgp/inline.py
+++ b/src/mailman_pgp/pgp/inline.py
@@ -24,23 +24,14 @@ from pgpy import PGPMessage
 from pgpy.constants import SymmetricKeyAlgorithm
 from public import public
 
-from mailman_pgp.utils.email import make_multipart
+from mailman_pgp.pgp.base import BaseWrapper
 from mailman_pgp.utils.pgp import key_from_blob, revoc_from_blob
 
 
 @public
-class InlineWrapper:
+class InlineWrapper(BaseWrapper):
     """Inline PGP wrapper."""
 
-    def __init__(self, msg):
-        """
-        Wrap the given message.
-
-        :param msg: The message to wrap.
-        :type msg: mailman.email.message.Message
-        """
-        self.msg = msg
-
     def get_payload(self):
         for part in walk(self.msg):
             if not part.is_multipart():
@@ -219,14 +210,21 @@ class InlineWrapper:
 
         :param key_revocations: A key revocation signature to attach.
         :type key_revocations: pgpy.PGPSignature
-        :return: The message with the signature attached.
-        :rtype: mailman.email.message.Message
+        :return:
+        :rtype: InlineWrapper
         """
-        out = make_multipart(self.msg)
+        if self.msg.get_content_type() != 'multipart/mixed':
+            # wrap in multipart/mixed
+            payload = copy.deepcopy(self.msg)
+            self.msg.set_payload([])
+            self.msg.set_type('multipart/mixed')
+            self.msg['MIME-Version'] = '1.0'
+            self.msg.attach(payload)
+
         for key_revocation in key_revocations:
             revoc_part = MIMEText(str(key_revocation))
-            out.attach(revoc_part)
-        return out
+            self.msg.attach(revoc_part)
+        return self
 
     def verify(self, key):
         """
@@ -250,11 +248,10 @@ class InlineWrapper:
 
         :param key: The key to sign with.
         :type key: pgpy.PGPKey
-        :return: The signed message.
-        :rtype: mailman.email.message.Message
+        :return:
+        :rtype: InlineWrapper
         """
-        out = copy.deepcopy(self.msg)
-        for part in walk(out):
+        for part in walk(self.msg):
             if not part.is_multipart():
                 if self._is_signed(part):
                     pmsg = PGPMessage.from_blob(part.get_payload())
@@ -263,16 +260,18 @@ class InlineWrapper:
                     pmsg = PGPMessage.new(payload, cleartext=True)
                 smsg = self._sign(pmsg, key, **kwargs)
                 part.set_payload(str(smsg))
-        return out
+        return self
 
     def _decrypt(self, part, key):
         message = PGPMessage.from_blob(part.get_payload())
-        # TODO: exception safe this.
         decrypted = key.decrypt(message)
         if decrypted.is_signed:
             part.set_payload(str(decrypted))
         else:
-            part.set_payload(decrypted.message)
+            dmsg = decrypted.message
+            if isinstance(dmsg, bytearray):
+                dmsg = dmsg.decode(decrypted.charset or 'utf-8')
+            part.set_payload(dmsg)
 
     def decrypt(self, key):
         """
@@ -280,14 +279,13 @@ class InlineWrapper:
 
         :param key: The key to decrypt with.
         :type key: pgpy.PGPKey
-        :return: The decrypted message.
-        :rtype: mailman.email.message.Message
+        :return:
+        :rtype: InlineWrapper
         """
-        out = copy.deepcopy(self.msg)
-        for part in walk(out):
+        for part in walk(self.msg):
             if not part.is_multipart() and self._is_encrypted(part):
                 self._decrypt(part, key)
-        return out
+        return self
 
     def _encrypt(self, pmsg, *keys, cipher, **kwargs):
         emsg = copy.copy(pmsg)
@@ -311,19 +309,19 @@ class InlineWrapper:
         :type keys: pgpy.PGPKey
         :param cipher: The symmetric cipher to use.
         :type cipher: SymmetricKeyAlgorithm
-        :return: mailman.email.message.Message
+        :return:
+        :rtype: InlineWrapper
         """
         if len(keys) == 0:
             raise ValueError('At least one key necessary.')
 
-        out = copy.deepcopy(self.msg)
-        for part in walk(out):
+        for part in walk(self.msg):
             if not part.is_multipart():
                 payload = str(part.get_payload())
                 pmsg = PGPMessage.new(payload)
                 emsg = self._encrypt(pmsg, *keys, cipher=cipher, **kwargs)
                 part.set_payload(str(emsg))
-        return out
+        return self
 
     def sign_encrypt(self, key, *keys, hash=None,
                      cipher=SymmetricKeyAlgorithm.AES256,
@@ -339,14 +337,13 @@ class InlineWrapper:
         :type hash: pgpy.constants.HashAlgorithm
         :param cipher:
         :type cipher: pgpy.constants.SymmetricKeyAlgorithm
-        :return: The signed + encrypted message.
-        :rtype: mailman.email.message.Message
+        :return:
+        :rtype: InlineWrapper
         """
         if len(keys) == 0:
             raise ValueError('At least one key necessary.')
 
-        out = copy.deepcopy(self.msg)
-        for part in walk(out):
+        for part in walk(self.msg):
             if not part.is_multipart():
                 if self._is_signed(part):
                     pmsg = PGPMessage.from_blob(part.get_payload())
@@ -356,10 +353,4 @@ class InlineWrapper:
                 smsg = self._sign(pmsg, key, hash=hash)
                 emsg = self._encrypt(smsg, *keys, cipher=cipher, **kwargs)
                 part.set_payload(str(emsg))
-        return out
-
-    def sign_then_encrypt(self, key, *keys, hash=None,
-                          cipher=SymmetricKeyAlgorithm.AES256,
-                          **kwargs):
-        return self.sign_encrypt(key, *keys, hash=hash, cipher=cipher,
-                                 **kwargs)
+        return self