feature/wrappers-modify

1 files changed, 91 insertions, 104 deletions

diff --git a/src/mailman_pgp/pgp/mime.py b/src/mailman_pgp/pgp/mime.py
index eeeed95..9ea6384 100644
--- a/src/mailman_pgp/pgp/mime.py
+++ b/src/mailman_pgp/pgp/mime.py
@@ -23,17 +23,18 @@ from email.iterators import walk
 from email.mime.application import MIMEApplication
 from email.utils import collapse_rfc2231_value
 
-from mailman.email.message import Message, MultipartDigestMessage
+from mailman.email.message import Message
 from pgpy import PGPMessage, PGPSignature
 from pgpy.constants import HashAlgorithm, SymmetricKeyAlgorithm
 from public import public
 
-from mailman_pgp.utils.email import copy_headers, make_multipart
+from mailman_pgp.pgp.base import BaseWrapper
+from mailman_pgp.utils.email import copy_headers
 from mailman_pgp.utils.pgp import key_from_blob, revoc_from_blob
 
 
 @public
-class MIMEWrapper:
+class MIMEWrapper(BaseWrapper):
     """PGP/MIME (RFC1847 + RFC3156) compliant wrapper."""
 
     _signature_subtype = 'pgp-signature'
@@ -44,20 +45,14 @@ class MIMEWrapper:
     _encrypted_type = 'application/' + _encryption_subtype
     _keys_type = 'application/' + _keys_subtype
 
+    _signed_multipart = 'multipart/signed'
+    _encrypted_multipart = 'multipart/encrypted'
+
     _signature_preamble = \
         'This is an OpenPGP/MIME signed message (RFC 4880 and 3156)'
     _encryption_preamble = \
         'This is an OpenPGP/MIME encrypted message (RFC 4880 and 3156)'
 
-    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):
         yield self.msg.as_string()
 
@@ -186,28 +181,40 @@ class MIMEWrapper:
                     continue
                 yield key
 
+    def _attach_key_part(self, obj, name, description):
+        key_part = MIMEApplication(_data=str(obj),
+                                   _subtype=MIMEWrapper._keys_subtype,
+                                   _encoder=encode_7or8bit,
+                                   name=name)
+        key_part.add_header('Content-Description', description)
+        key_part.add_header('Content-Disposition', 'attachment',
+                            filename=name)
+        self.msg.attach(key_part)
+
     def attach_keys(self, *keys):
         """
         Attach a key to this message, as per RFC3156 section 7.
 
-        :param keys: A key to attach.
+        :param keys: Keys to attach.
         :type keys: pgpy.PGPKey
-        :return: The message with the key attached.
-        :rtype: mailman.email.message.Message
+        :return:
+        :rtype: MIMEWrapper
         """
-        out = make_multipart(self.msg)
+        if len(keys) == 0:
+            return self
+
+        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 in keys:
             filename = '0x' + key.fingerprint.keyid + '.asc'
-            key_part = MIMEApplication(_data=str(key),
-                                       _subtype=MIMEWrapper._keys_subtype,
-                                       _encoder=encode_7or8bit,
-                                       name=filename)
-            key_part.add_header('Content-Description',
-                                'OpenPGP key')
-            key_part.add_header('Content-Disposition', 'attachment',
-                                filename=filename)
-            out.attach(key_part)
-        return out
+            self._attach_key_part(key, filename, 'OpenPGP key')
+        return self
 
     def _is_revoc(self, part):
         if part.get_content_type() != MIMEWrapper._keys_type:
@@ -246,22 +253,25 @@ class MIMEWrapper:
 
         :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: MIMEWrapper
         """
-        out = make_multipart(self.msg)
+        if len(key_revocations) == 0:
+            return self
+
+        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:
             filename = '0x' + key_revocation.signer + '.asc'
-            revoc_part = MIMEApplication(_data=str(key_revocation),
-                                         _subtype=MIMEWrapper._keys_subtype,
-                                         _encoder=encode_7or8bit,
-                                         name=filename)
-            revoc_part.add_header('Content-Description',
-                                  'OpenPGP key')
-            revoc_part.add_header('Content-Disposition', 'attachment',
-                                  filename=filename)
-            out.attach(revoc_part)
-        return out
+            self._attach_key_part(key_revocation, filename,
+                                  'OpenPGP key revocation')
+        return self
 
     def verify(self, key):
         """
@@ -298,12 +308,13 @@ class MIMEWrapper:
 
         :param msg:
         :param signature:
-        :return:
         """
-        micalg = self._micalg(signature.hash_algorithm)
-        out = MultipartDigestMessage('signed', micalg=micalg,
-                                     protocol=MIMEWrapper._signed_type)
-        out.preamble = MIMEWrapper._signature_preamble
+        self.msg.set_payload([])
+        self.msg.attach(msg)
+        self.msg.set_type(MIMEWrapper._signed_multipart)
+        self.msg.set_param('micalg', self._micalg(signature.hash_algorithm))
+        self.msg.set_param('protocol', MIMEWrapper._signed_type)
+        self.msg.preamble = MIMEWrapper._signature_preamble
 
         second_part = MIMEApplication(_data=str(signature),
                                       _subtype=MIMEWrapper._signature_subtype,
@@ -313,11 +324,7 @@ class MIMEWrapper:
                                'OpenPGP digital signature')
         second_part.add_header('Content-Disposition', 'attachment',
                                filename='signature.asc')
-
-        out.attach(copy.deepcopy(msg))
-        out.attach(second_part)
-        copy_headers(msg, out)
-        return out
+        self.msg.attach(second_part)
 
     def sign(self, key, **kwargs):
         """
@@ -325,12 +332,14 @@ class MIMEWrapper:
 
         :param key: The key to sign with.
         :type key: pgpy.PGPKey
-        :return: The signed message.
-        :rtype: mailman.email.message.Message
+        :rtype: MIMEWrapper
         """
         payload = next(iter(self.get_payload()))
         signature = key.sign(payload, **kwargs)
-        return self._wrap_signed(self.msg, signature)
+        original_msg = copy.deepcopy(self.msg)
+        self._wrap_signed(original_msg, signature)
+
+        return self
 
     def decrypt(self, key):
         """
@@ -342,7 +351,6 @@ class MIMEWrapper:
         :rtype: mailman.email.message.Message
         """
         pmsg = next(iter(self.get_encrypted()))
-        # TODO: exception safe this.
         decrypted = key.decrypt(pmsg)
 
         dmsg = decrypted.message
@@ -351,9 +359,17 @@ class MIMEWrapper:
 
         out = message_from_string(dmsg, _class=Message)
         if decrypted.is_signed:
-            out = self._wrap_signed(out, decrypted.signatures.pop())
-        copy_headers(self.msg, out)
-        return out
+            # rewrap, self.msg should be multipart/signed,
+            # headers from out should overwrite those from self.msg
+            # self.msg payload should be [out, sig]
+            signature = next(iter(decrypted.signatures))
+            self._wrap_signed(out, signature)
+        else:
+            # self.msg payload should be out.get_payload
+            # headers from out should overwrite those from self.msg
+            self.msg.set_payload(out.get_payload())
+            copy_headers(out, self.msg, True)
+        return self
 
     def _encrypt(self, pmsg, *keys, cipher, **kwargs):
         emsg = copy.copy(pmsg)
@@ -369,16 +385,16 @@ class MIMEWrapper:
         return emsg
 
     def _wrap_encrypted(self, payload):
-        out = MultipartDigestMessage('encrypted',
-                                     protocol=MIMEWrapper._encrypted_type)
-        out.preamble = MIMEWrapper._encryption_preamble
-
+        self.msg.set_payload([])
+        self.msg.set_type(MIMEWrapper._encrypted_multipart)
+        self.msg.set_param('protocol', MIMEWrapper._encrypted_type)
+        self.msg.preamble = MIMEWrapper._encryption_preamble
         first_part = MIMEApplication(_data='Version: 1',
                                      _subtype=MIMEWrapper._encryption_subtype,
                                      _encoder=encode_7or8bit)
         first_part.add_header('Content-Description',
                               'PGP/MIME version identification')
-
+        self.msg.attach(first_part)
         second_part = MIMEApplication(_data=str(payload),
                                       _subtype='octet-stream',
                                       _encoder=encode_7or8bit,
@@ -387,9 +403,7 @@ class MIMEWrapper:
                                'OpenPGP encrypted message')
         second_part.add_header('Content-Disposition', 'inline',
                                filename='encrypted.asc')
-        out.attach(first_part)
-        out.attach(second_part)
-        return out
+        self.msg.attach(second_part)
 
     def encrypt(self, *keys, cipher=SymmetricKeyAlgorithm.AES256,
                 **kwargs):
@@ -400,21 +414,25 @@ class MIMEWrapper:
         :type keys: pgpy.PGPKey
         :param cipher: The symmetric cipher to use.
         :type cipher: pgpy.constants.SymmetricKeyAlgorithm
-        :return: The encrypted message.
-        :rtype: mailman.email.message.Message
+        :return:
+        :rtype: MIMEWrapper
         """
         if len(keys) == 0:
             raise ValueError('At least one key necessary.')
 
         if self.is_signed():
+            # self.msg payload should be [ version_1, encrypted]
+            # headers should remain the same, except Content-Type
+            # signature should be combined into the PGP blob
             pmsg = PGPMessage.new(next(iter(self.get_signed())))
             pmsg |= next(iter(self.get_signature()))
         else:
+            # self.msg payload should be [ version_1, encrypted]
+            # headers should remain the same, except Content-Type
             pmsg = PGPMessage.new(next(iter(self.get_payload())))
         pmsg = self._encrypt(pmsg, *keys, cipher=cipher, **kwargs)
-        out = self._wrap_encrypted(pmsg)
-        copy_headers(self.msg, out)
-        return out
+        self._wrap_encrypted(pmsg)
+        return self
 
     def sign_encrypt(self, key, *keys, hash=None,
                      cipher=SymmetricKeyAlgorithm.AES256,
@@ -432,8 +450,8 @@ class MIMEWrapper:
         :type hash: pgpy.constants.HashAlgorithm
         :param cipher:
         :type cipher: pgpy.constants.SymmetricKeyAlgorithm
-        :return: The signed + encrypted message.
-        :rtype: mailman.email.message.Message
+        :return:
+        :rtype: MIMEWrapper
         """
         if len(keys) == 0:
             raise ValueError('At least one key necessary.')
@@ -442,36 +460,5 @@ class MIMEWrapper:
         pmsg = PGPMessage.new(payload)
         pmsg |= key.sign(pmsg, hash=hash)
         pmsg = self._encrypt(pmsg, *keys, cipher=cipher, **kwargs)
-        out = self._wrap_encrypted(pmsg)
-        copy_headers(self.msg, out)
-        return out
-
-    def sign_then_encrypt(self, key, *keys, hash=None,
-                          cipher=SymmetricKeyAlgorithm.AES256,
-                          **kwargs):
-        """
-        Sign then encrypt the message.
-
-        This is as per RFC 3156 section 6.1 - RFC 1847 Encapsulation.
-
-        :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
-        """
-        if len(keys) == 0:
-            raise ValueError('At least one key necessary.')
-
-        out = self.sign(key, hash=hash)
-        out_wrapped = MIMEWrapper(out)
-        pmsg = PGPMessage.new(next(out_wrapped.get_payload()))
-        pmsg = self._encrypt(pmsg, *keys, cipher=cipher, **kwargs)
-        out = self._wrap_encrypted(pmsg)
-        copy_headers(self.msg, out)
-        return out
+        self._wrap_encrypted(pmsg)
+        return self