1 files changed, 72 insertions, 17 deletions

diff --git a/src/mailman_pgp/pgp/mime.py b/src/mailman_pgp/pgp/mime.py
index a1303c9..03177ab 100644
--- a/src/mailman_pgp/pgp/mime.py
+++ b/src/mailman_pgp/pgp/mime.py
@@ -28,8 +28,8 @@ from pgpy import PGPDetachedSignature, PGPMessage
 from pgpy.constants import HashAlgorithm, SymmetricKeyAlgorithm
 from public import public
 
-from mailman_pgp.utils.email import copy_headers
-from mailman_pgp.utils.pgp import key_from_blob
+from mailman_pgp.utils.email import copy_headers, make_multipart
+from mailman_pgp.utils.pgp import key_from_blob, revoc_from_blob
 
 
 @public
@@ -63,7 +63,7 @@ class MIMEWrapper:
 
     def _is_mime(self):
         is_multipart = self.msg.is_multipart()
-        payloads = len(self.msg.get_payload())
+        payloads = len(self.msg.get_payload()) if self.msg.get_payload() else 0
 
         return is_multipart and payloads == 2
 
@@ -186,26 +186,81 @@ class MIMEWrapper:
                     continue
                 yield key
 
-    def attach_key(self, key):
+    def attach_keys(self, *keys):
         """
         Attach a key to this message, as per RFC3156 section 7.
 
-        :param key: A key to attach.
-        :type key: pgpy.PGPKey
+        :param keys: A key to attach.
+        :type keys: pgpy.PGPKey
         :return: The message with the key attached.
         :rtype: mailman.email.message.Message
         """
-        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 = copy.deepcopy(self.msg)
-        out.attach(key_part)
+        out = make_multipart(self.msg)
+        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
+
+    def _is_revoc(self, part):
+        if part.get_content_type() != MIMEWrapper._keys_type:
+            return False
+        try:
+            revoc_from_blob(part.get_payload())
+        except ValueError:
+            return False
+        return True
+
+    def is_revocs(self):
+        for part in walk(self.msg):
+            if (not part.is_multipart() and not self._is_revoc(part)):
+                return False
+        return True
+
+    def has_revocs(self):
+        for part in walk(self.msg):
+            if (not part.is_multipart() and self._is_revoc(part)):
+                return True
+        return False
+
+    def revocs(self):
+        for part in walk(self.msg):
+            if (not part.is_multipart()  # noqa
+                and part.get_content_type() == MIMEWrapper._keys_type):
+                try:
+                    revoc = revoc_from_blob(part.get_payload())
+                except:
+                    continue
+                yield revoc
+
+    def attach_revocs(self, *key_revocations):
+        """
+        Attach a key revocation signature to the message, as a key subpart.
+
+        :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
+        """
+        out = make_multipart(self.msg)
+        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
 
     def verify(self, key):