2 files changed, 81 insertions, 23 deletions

diff --git a/src/mailman_pgp/pgp/inline.py b/src/mailman_pgp/pgp/inline.py
index fe5b2c3..ca83e89 100644
--- a/src/mailman_pgp/pgp/inline.py
+++ b/src/mailman_pgp/pgp/inline.py
@@ -159,11 +159,10 @@ class InlineWrapper:
             if not part.is_multipart() and self._is_signed(part):
                 yield self._verify(part, key)
 
-    def _sign(self, part, key, hash):
-        payload = str(part.get_payload())
-        pmsg = PGPMessage.new(payload, cleartext=True)
-        pmsg |= key.sign(pmsg, hash=hash)
-        part.set_payload(str(pmsg))
+    def _sign(self, pmsg, key, hash):
+        smsg = copy.copy(pmsg)
+        smsg |= key.sign(smsg, hash=hash)
+        return smsg
 
     def sign(self, key, hash=None):
         """
@@ -179,7 +178,10 @@ class InlineWrapper:
         out = copy.deepcopy(self.msg)
         for part in walk(out):
             if not part.is_multipart():
-                self._sign(part, key, hash)
+                payload = str(part.get_payload())
+                pmsg = PGPMessage.new(payload, cleartext=True)
+                smsg = self._sign(pmsg, key, hash)
+                part.set_payload(str(smsg))
         return out
 
     def _decrypt(self, part, key):
@@ -202,12 +204,17 @@ class InlineWrapper:
                 self._decrypt(part, key)
         return out
 
-    def _encrypt(self, part, *keys, **kwargs):
-        payload = str(part.get_payload())
-        pmsg = PGPMessage.new(payload)
-        for key in keys:
-            pmsg = key.encrypt(pmsg, **kwargs)
-        part.set_payload(str(pmsg))
+    def _encrypt(self, pmsg, *keys, cipher):
+        emsg = copy.copy(pmsg)
+        if len(keys) == 1:
+            emsg = keys[0].encrypt(emsg, cipher=cipher)
+        else:
+            session_key = cipher.gen_key()
+            for key in keys:
+                emsg = key.encrypt(emsg, cipher=cipher,
+                                   session_key=session_key)
+            del session_key
+        return emsg
 
     def encrypt(self, *keys, cipher=SymmetricKeyAlgorithm.AES256):
         """
@@ -225,12 +232,35 @@ class InlineWrapper:
         out = copy.deepcopy(self.msg)
         for part in walk(out):
             if not part.is_multipart():
-                if len(keys) == 1:
-                    self._encrypt(part, *keys, cipher=cipher)
-                else:
-                    session_key = cipher.gen_key()
-                    for key in keys:
-                        self._encrypt(part, key, cipher=cipher,
-                                      session_key=session_key)
-                    del session_key
+                payload = str(part.get_payload())
+                pmsg = PGPMessage.new(payload)
+                emsg = self._encrypt(pmsg, *keys, cipher=cipher)
+                part.set_payload(str(emsg))
         return out
+
+    def sign_encrypt(self, key, *keys, hash=None,
+                     cipher=SymmetricKeyAlgorithm.AES256):
+        """
+
+        :param key:
+        :param keys:
+        :param hash:
+        :param cipher:
+        :return:
+        """
+        if len(keys) == 0:
+            raise ValueError('At least one key necessary.')
+
+        out = copy.deepcopy(self.msg)
+        for part in walk(out):
+            if not part.is_multipart():
+                payload = str(part.get_payload())
+                pmsg = PGPMessage.new(payload)
+                smsg = self._sign(pmsg, key, hash=hash)
+                emsg = self._encrypt(smsg, *keys, cipher=cipher)
+                part.set_payload(str(emsg))
+        return out
+
+    def sign_then_encrypt(self, key, *keys, hash=None,
+                          cipher=SymmetricKeyAlgorithm.AES256):
+        return self.sign_encrypt(key, *keys, hash=hash, cipher=cipher)
\ No newline at end of file
diff --git a/src/mailman_pgp/pgp/mime.py b/src/mailman_pgp/pgp/mime.py
index b28fc54..47ed192 100644
--- a/src/mailman_pgp/pgp/mime.py
+++ b/src/mailman_pgp/pgp/mime.py
@@ -266,6 +266,9 @@ class MIMEWrapper:
         :return: The encrypted message.
         :rtype: mailman.email.message.Message
         """
+        if len(keys) == 0:
+            raise ValueError('At least one key necessary.')
+
         payload = self.msg.as_string()
         pmsg = PGPMessage.new(payload)
         pmsg = self._encrypt(pmsg, *keys, cipher=cipher)
@@ -274,17 +277,24 @@ class MIMEWrapper:
     def sign_encrypt(self, key, *keys, hash=None,
                      cipher=SymmetricKeyAlgorithm.AES256):
         """
-        Sign and encrypt te message, in one go.
+        Sign and encrypt the message, in one go.
+
+        This is as per RFC 3156 section 6.2 - Combined method.
 
         :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.')
+
         payload = self.msg.as_string()
         pmsg = PGPMessage.new(payload)
         pmsg |= key.sign(pmsg, hash=hash)
@@ -293,7 +303,25 @@ class MIMEWrapper:
 
     def sign_then_encrypt(self, key, *keys, hash=None,
                           cipher=SymmetricKeyAlgorithm.AES256):
+        """
+        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)
         out_wrapped = MIMEWrapper(out)
-        out = out_wrapped.encrypt(*keys, cipher=cipher)
-        return out
+        return out_wrapped.encrypt(*keys, cipher=cipher)