From 22a688b4f0579b1e1e51b4164934fe2afc357671 Mon Sep 17 00:00:00 2001
From: J08nY
Date: Wed, 9 Aug 2017 17:07:10 +0200
Subject: Make MIMEWrapper encrypt signed messages by combining.

---
 src/mailman_pgp/mta/bulk.py            |  2 ++
 src/mailman_pgp/mta/personalized.py    |  4 ++++
 src/mailman_pgp/pgp/mime.py            | 22 +++++++++++++++-------
 src/mailman_pgp/pgp/tests/test_mime.py |  4 +++-
 4 files changed, 24 insertions(+), 8 deletions(-)

(limited to 'src')

diff --git a/src/mailman_pgp/mta/bulk.py b/src/mailman_pgp/mta/bulk.py
index a2cc8c6..a01b5a7 100644
--- a/src/mailman_pgp/mta/bulk.py
+++ b/src/mailman_pgp/mta/bulk.py
@@ -97,6 +97,8 @@ class PGPBulkMixin:
             else:
                 out = wrapped.sign(pgp_list.key)
         else:
+            # Definitely encrypt here, the case where we don't encrypt or sign
+            # is handled above at the start of the func.
             out = wrapped.encrypt(pgp_list.pubkey, *keys, throw_keyid=True)
 
         overwrite_message(out, msg)
diff --git a/src/mailman_pgp/mta/personalized.py b/src/mailman_pgp/mta/personalized.py
index a89301f..bd50c70 100644
--- a/src/mailman_pgp/mta/personalized.py
+++ b/src/mailman_pgp/mta/personalized.py
@@ -16,6 +16,8 @@
 # this program.  If not, see <http://www.gnu.org/licenses/>.
 
 """PGP enabled IndividualDelivery."""
+import copy
+
 from mailman.mta.base import IndividualDelivery
 from mailman.mta.decorating import DecoratingMixin
 from mailman.mta.personalized import PersonalizedMixin
@@ -67,6 +69,8 @@ class PGPIndividualMixin:
             else:
                 out = wrapped.sign(pgp_list.key)
         else:
+            # Definitely encrypt here, the case where we don't encrypt or sign
+            # is handled above at the start of the func.
             out = wrapped.encrypt(key, pgp_list.pubkey)
 
         overwrite_message(out, msg)
diff --git a/src/mailman_pgp/pgp/mime.py b/src/mailman_pgp/pgp/mime.py
index 03177ab..32e2cab 100644
--- a/src/mailman_pgp/pgp/mime.py
+++ b/src/mailman_pgp/pgp/mime.py
@@ -24,7 +24,7 @@ from email.mime.application import MIMEApplication
 from email.utils import collapse_rfc2231_value
 
 from mailman.email.message import Message, MultipartDigestMessage
-from pgpy import PGPDetachedSignature, PGPMessage
+from pgpy import PGPMessage, PGPSignature, PGPDetachedSignature
 from pgpy.constants import HashAlgorithm, SymmetricKeyAlgorithm
 from public import public
 
@@ -358,16 +358,17 @@ class MIMEWrapper:
         return out
 
     def _encrypt(self, pmsg, *keys, cipher, **kwargs):
+        emsg = copy.copy(pmsg)
         if len(keys) == 1:
-            pmsg = keys[0].encrypt(pmsg, cipher=cipher, **kwargs)
+            emsg = keys[0].encrypt(emsg, cipher=cipher, **kwargs)
         else:
             session_key = cipher.gen_key()
             for key in keys:
-                pmsg = key.encrypt(pmsg, cipher=cipher,
+                emsg = key.encrypt(emsg, cipher=cipher,
                                    sessionkey=session_key,
                                    **kwargs)
             del session_key
-        return pmsg
+        return emsg
 
     def _wrap_encrypted(self, payload):
         out = MultipartDigestMessage('encrypted',
@@ -407,8 +408,11 @@ class MIMEWrapper:
         if len(keys) == 0:
             raise ValueError('At least one key necessary.')
 
-        payload = self.msg.as_string()
-        pmsg = PGPMessage.new(payload)
+        if self.is_signed():
+            pmsg = PGPMessage.new(next(iter(self.get_signed())))
+            pmsg |= next(iter(self.get_signature()))
+        else:
+            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)
@@ -468,4 +472,8 @@ class MIMEWrapper:
 
         out = self.sign(key, hash)
         out_wrapped = MIMEWrapper(out)
-        return out_wrapped.encrypt(*keys, cipher=cipher, **kwargs)
+        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
diff --git a/src/mailman_pgp/pgp/tests/test_mime.py b/src/mailman_pgp/pgp/tests/test_mime.py
index e9951cf..e599170 100644
--- a/src/mailman_pgp/pgp/tests/test_mime.py
+++ b/src/mailman_pgp/pgp/tests/test_mime.py
@@ -111,7 +111,9 @@ class TestEncryption(MIMEWrapperTestCase):
          load_key('rsa_1024.pub.asc')),
         (load_message('clear_multipart.eml'),
          (load_key('rsa_1024.pub.asc'),
-          load_key('ecc_p256.pub.asc')))
+          load_key('ecc_p256.pub.asc'))),
+        (load_message('mime_signed.eml'),
+         load_key('ecc_p256.pub.asc'))
     ])
     def test_encrypt(self, message, keys, **kwargs):
         if isinstance(keys, tuple):
-- 
cgit v1.2.3-70-g09d2