6 files changed, 153 insertions, 104 deletions

diff --git a/src/mailman_pgp/commands/eml_key.py b/src/mailman_pgp/commands/eml_key.py
index 9df6065..7b7782d 100644
--- a/src/mailman_pgp/commands/eml_key.py
+++ b/src/mailman_pgp/commands/eml_key.py
@@ -19,6 +19,7 @@
 from email.utils import parseaddr
 
 from mailman.interfaces.command import ContinueProcessing, IEmailCommand
+from mailman.interfaces.pending import IPendings
 from mailman.interfaces.subscriptions import ISubscriptionManager
 from mailman.interfaces.usermanager import IUserManager
 from public import public
@@ -47,6 +48,10 @@ def _cmd_set(pgp_list, mlist, msg, msgdata, arguments, results):
         return ContinueProcessing.no
 
     wrapped = PGPWrapper(msg)
+    if wrapped.is_encrypted():
+        decrypted = wrapped.decrypt(pgp_list.key)
+        wrapped = PGPWrapper(decrypted)
+
     if not wrapped.has_keys():
         print('No keys attached? Send a key.', file=results)
         return ContinueProcessing.no
@@ -66,21 +71,24 @@ def _cmd_set(pgp_list, mlist, msg, msgdata, arguments, results):
         print('No adddress to subscribe with.', file=results)
         return ContinueProcessing.no
 
-    with transaction() as t:
-        pgp_address = PGPAddress.for_address(address)
-        if pgp_address is None:
-            pgp_address = PGPAddress(address)
-        pgp_address.key = keys.pop()
-        t.add(pgp_address)
+    pgp_address = PGPAddress.for_address(address)
+    if pgp_address is None:
+        print('A pgp enabled address not found.', file=results)
+        return ContinueProcessing.no
 
     token = arguments[1]
-    try:
-        ISubscriptionManager(mlist).confirm(token)
-        print('Key succesfully set.', file=results)
-        print('Key fingerprint: {}'.format(pgp_address.key.fingerprint),
-              file=results)
-    except LookupError:
+    pendable = getUtility(IPendings).confirm(token, expunge=False)
+    if pendable is None:
         print('Wrong token.', file=results)
+        return ContinueProcessing.no
+
+    with transaction():
+        pgp_address.key = keys.pop()
+    ISubscriptionManager(mlist).confirm(token)
+
+    print('Key succesfully set.', file=results)
+    print('Key fingerprint: {}'.format(pgp_address.key.fingerprint),
+          file=results)
 
     return ContinueProcessing.no
 
@@ -115,13 +123,17 @@ def _cmd_confirm(pgp_list, mlist, msg, msgdata, arguments, results):
 
     token = arguments[1]
 
-    expecting = CONFIRM_REQUEST.format(pgp_address.key_fingerprint,
-                                       token)
+    pendable = getUtility(IPendings).confirm(token, expunge=False)
+    if pendable is None:
+        print('Wrong token.', file=results)
+        return ContinueProcessing.no
+
+    # TODO differentiate between key change and subscription here.
+
+    expecting = CONFIRM_REQUEST.format(pgp_address.key_fingerprint, token)
     for sig_subject in wrapped.get_signed():
         if expecting in sig_subject:
-            with transaction():
-                pgp_address.key_confirmed = True
-                ISubscriptionManager(mlist).confirm(token)
+            ISubscriptionManager(mlist).confirm(token)
             break
     else:
         print("Message doesn't contain the expected statement.", file=results)
diff --git a/src/mailman_pgp/commands/tests/test_key.py b/src/mailman_pgp/commands/tests/test_key.py
index fe75a6a..64f8ae6 100644
--- a/src/mailman_pgp/commands/tests/test_key.py
+++ b/src/mailman_pgp/commands/tests/test_key.py
@@ -134,7 +134,6 @@ class TestPreSubscription(unittest.TestCase):
         pgp_address = PGPAddress.for_address(bart)
         self.assertIsNotNone(pgp_address)
         self.assertEqual(pgp_address.key.fingerprint, bart_key.fingerprint)
-        self.assertEqual(pgp_address.key_fingerprint, bart_key.fingerprint)
         self.assertFalse(pgp_address.key_confirmed)
 
         items = get_queue_messages('virgin', expected_count=2)
@@ -242,11 +241,6 @@ class TestPreSubscription(unittest.TestCase):
 
         get_queue_messages('virgin')
 
-        with transaction() as t:
-            pgp_address = PGPAddress(bart)
-            pgp_address.key = bart_key.pubkey
-            t.add(pgp_address)
-
         message = _create_plain('bart@example.com', 'test@example.com',
                                 'Re: key confirm {}'.format(token),
                                 CONFIRM_REQUEST.format(bart_key.fingerprint,
@@ -275,11 +269,6 @@ class TestPreSubscription(unittest.TestCase):
 
         get_queue_messages('virgin')
 
-        with transaction() as t:
-            pgp_address = PGPAddress(bart)
-            pgp_address.key = bart_key.pubkey
-            t.add(pgp_address)
-
         message = _create_plain('bart@example.com', 'test@example.com',
                                 'Re: key confirm {}'.format(token),
                                 CONFIRM_REQUEST.format(bart_key.fingerprint,
@@ -290,6 +279,7 @@ class TestPreSubscription(unittest.TestCase):
 
         mm_config.switchboards['command'].enqueue(message,
                                                   listid='test.example.com')
+
         make_testable_runner(CommandRunner, 'command').run()
 
         pgp_address = PGPAddress.for_address(bart)
@@ -346,11 +336,6 @@ class TestPreSubscription(unittest.TestCase):
 
         get_queue_messages('virgin')
 
-        with transaction() as t:
-            pgp_address = PGPAddress(bart)
-            pgp_address.key = bart_key.pubkey
-            t.add(pgp_address)
-
         message = _create_plain('bart@example.com', 'test@example.com',
                                 'Re: key confirm {}'.format(token),
                                 CONFIRM_REQUEST.format(bart_key.fingerprint,
@@ -408,3 +393,35 @@ class TestAfterSubscription(unittest.TestCase):
 
         confirm_wrapped = PGPWrapper(confirm_request)
         self.assertTrue(confirm_wrapped.is_encrypted())
+        decrypted = confirm_wrapped.decrypt(bart_new_key)
+        self.assertIn('key confirm', decrypted['subject'])
+
+    def test_key_change_confirm(self):
+        bart = getUtility(IUserManager).create_address('bart@example.com',
+                                                       'Bart Person')
+        bart_key = load_key('rsa_1024.priv.asc')
+        bart_new_key = load_key('ecc_p256.priv.asc')
+
+        with transaction() as t:
+            pgp_address = PGPAddress(bart)
+            pgp_address.key = bart_key.pubkey
+            pgp_address.key_confirmed = True
+            t.add(pgp_address)
+
+        message = _create_mixed('bart@example.com', 'test@example.com',
+                                'key change')
+        wrapped_message = MIMEWrapper(message)
+        message = wrapped_message.attach_key(bart_new_key.pubkey)
+
+        mm_config.switchboards['command'].enqueue(message,
+                                                  listid='test.example.com')
+        make_testable_runner(CommandRunner, 'command').run()
+
+        items = get_queue_messages('virgin', expected_count=2)
+        if items[0].msg['Subject'] == 'The results of your email commands':
+            confirm_request = items[1].msg
+        else:
+            confirm_request = items[0].msg
+        request_wrapped = PGPWrapper(confirm_request)
+        request_wrapped.decrypt(bart_new_key)
+        # TODO finish this
diff --git a/src/mailman_pgp/workflows/base.py b/src/mailman_pgp/workflows/base.py
index a8679bd..d05781d 100644
--- a/src/mailman_pgp/workflows/base.py
+++ b/src/mailman_pgp/workflows/base.py
@@ -20,6 +20,7 @@ from mailman.email.message import UserNotification
 from mailman.interfaces.subscriptions import TokenOwner
 from pgpy import PGPKey
 
+from mailman_pgp.database import transaction
 from mailman_pgp.model.address import PGPAddress
 from mailman_pgp.model.list import PGPMailingList
 from mailman_pgp.pgp.utils import copy_headers
@@ -44,6 +45,15 @@ Token: {}
 """
 
 
+class PGPMixin:
+    def _step_pgp_prepare(self):
+        pgp_address = PGPAddress.for_address(self.address)
+        if pgp_address is None:
+            with transaction() as t:
+                pgp_address = PGPAddress(self.address)
+                t.add(pgp_address)
+
+
 class SetPubkeyMixin:
     def __init__(self, pubkey=None, pre_confirmed=False):
         self.pubkey = pubkey
@@ -64,13 +74,14 @@ class SetPubkeyMixin:
 
     def _step_pubkey_checks(self):
         pgp_address = PGPAddress.for_address(self.address)
+        assert pgp_address is not None
 
-        if pgp_address is not None:
-            if not pgp_address.key:
+        if self.pubkey is None:
+            if pgp_address.key is None:
                 self.push('send_key_request')
         else:
-            if not self.pubkey:
-                self.push('send_key_request')
+            with transaction():
+                pgp_address.key = self.pubkey
 
     def _step_send_key_request(self):
         self._set_token(TokenOwner.subscriber)
@@ -89,14 +100,6 @@ class SetPubkeyMixin:
         self._restore_subscriber()
         self._set_token(TokenOwner.no_one)
 
-        pgp_address = PGPAddress.for_address(self.address)
-        if pgp_address is None or pgp_address.key is None:
-            # The workflow was confirmed but we still dont have an address
-            # or the pubkey. So resend request and wait.
-            self.push('send_key_request')
-        else:
-            self.pubkey = pgp_address.key
-
 
 class ConfirmPubkeyMixin:
     def __init__(self, pre_confirmed=False):
@@ -104,27 +107,31 @@ class ConfirmPubkeyMixin:
 
     def _step_pubkey_confirmation(self):
         pgp_address = PGPAddress.for_address(self.address)
+        assert pgp_address is not None
 
-        if pgp_address is not None:
-            if not pgp_address.key_confirmed and not self.pubkey_confirmed:
-                self.push('send_key_confirm_request')
+        if self.pubkey_confirmed:
+            with transaction():
+                pgp_address.key_confirmed = True
         else:
-            if not self.pubkey_confirmed:
+            if not pgp_address.key_confirmed:
                 self.push('send_key_confirm_request')
 
     def _step_send_key_confirm_request(self):
         self._set_token(TokenOwner.subscriber)
         self.push('receive_key_confirmation')
         self.save()
+
+        pgp_address = PGPAddress.for_address(self.address)
         request_address = self.mlist.request_address
         email_address = self.address.email
         msg = UserNotification(email_address, request_address,
                                'key confirm {}'.format(self.token),
-                               CONFIRM_REQUEST.format(self.pubkey.fingerprint,
-                                                      self.token))
+                               CONFIRM_REQUEST.format(
+                                       pgp_address.key_fingerprint,
+                                       self.token))
         pgp_list = PGPMailingList.for_list(self.mlist)
         wrapped = PGPWrapper(msg)
-        encrypted = wrapped.sign_encrypt(pgp_list.key, self.pubkey,
+        encrypted = wrapped.sign_encrypt(pgp_list.key, pgp_address.key,
                                          pgp_list.pubkey)
 
         msg.set_payload(encrypted.get_payload())
@@ -135,3 +142,6 @@ class ConfirmPubkeyMixin:
     def _step_receive_key_confirmation(self):
         self._restore_subscriber()
         self._set_token(TokenOwner.no_one)
+        with transaction():
+            pgp_address = PGPAddress.for_address(self.address)
+            pgp_address.key_confirmed = True
diff --git a/src/mailman_pgp/workflows/key_change.py b/src/mailman_pgp/workflows/key_change.py
index 0098b19..8831d28 100644
--- a/src/mailman_pgp/workflows/key_change.py
+++ b/src/mailman_pgp/workflows/key_change.py
@@ -95,8 +95,8 @@ class KeyChangeWorkflow(Workflow):
         msg = UserNotification(email_address, request_address,
                                'key confirm {}'.format(self.token),
                                CHANGE_CONFIRM_REQUEST.format(
-                                   self.pubkey.fingerprint,
-                                   self.token))
+                                       self.pubkey.fingerprint,
+                                       self.token))
         wrapped = PGPWrapper(msg)
         encrypted = wrapped.encrypt(self.pubkey)
 
@@ -106,7 +106,13 @@ class KeyChangeWorkflow(Workflow):
         raise StopIteration
 
     def _step_receive_confirmation(self):
-        pass
+        self.pgp_address.key = self.pubkey
+        self.pgp_address.key_confirmed = True
+
+        pendings = getUtility(IPendings)
+        if self.token is not None:
+            pendings.confirm(self.token)
+            self.token = None
 
     @classmethod
     def pendable_class(cls):
diff --git a/src/mailman_pgp/workflows/subscription.py b/src/mailman_pgp/workflows/subscription.py
index c4138e6..ce02013 100644
--- a/src/mailman_pgp/workflows/subscription.py
+++ b/src/mailman_pgp/workflows/subscription.py
@@ -24,13 +24,15 @@ from mailman.workflows.common import (ConfirmationMixin, ModerationMixin,
 from public import public
 from zope.interface import implementer
 
-from mailman_pgp.workflows.base import ConfirmPubkeyMixin, SetPubkeyMixin
+from mailman_pgp.workflows.base import (ConfirmPubkeyMixin, PGPMixin,
+                                        SetPubkeyMixin)
 
 
 @public
 @implementer(ISubscriptionWorkflow)
 class OpenSubscriptionPolicy(SubscriptionBase, VerificationMixin,
-                             SetPubkeyMixin, ConfirmPubkeyMixin):
+                             SetPubkeyMixin, ConfirmPubkeyMixin,
+                             PGPMixin):
     """"""
 
     name = 'pgp-policy-open'
@@ -54,11 +56,13 @@ class OpenSubscriptionPolicy(SubscriptionBase, VerificationMixin,
         VerificationMixin.__init__(self, pre_verified=pre_verified)
         SetPubkeyMixin.__init__(self, pubkey=pubkey)
         ConfirmPubkeyMixin.__init__(self, pre_confirmed=pubkey_pre_confirmed)
+        PGPMixin.__init__(self)
 
     def _step_prepare(self):
         self.push('do_subscription')
         self.push('pubkey_confirmation')
         self.push('pubkey_checks')
+        self.push('pgp_prepare')
         self.push('verification_checks')
         self.push('sanity_checks')
 
@@ -67,7 +71,7 @@ class OpenSubscriptionPolicy(SubscriptionBase, VerificationMixin,
 @implementer(ISubscriptionWorkflow)
 class ConfirmSubscriptionPolicy(SubscriptionBase, VerificationMixin,
                                 ConfirmationMixin, SetPubkeyMixin,
-                                ConfirmPubkeyMixin):
+                                ConfirmPubkeyMixin, PGPMixin):
     """"""
 
     name = 'pgp-policy-confirm'
@@ -93,11 +97,13 @@ class ConfirmSubscriptionPolicy(SubscriptionBase, VerificationMixin,
         ConfirmationMixin.__init__(self, pre_confirmed=pre_confirmed)
         SetPubkeyMixin.__init__(self, pubkey=pubkey)
         ConfirmPubkeyMixin.__init__(self, pre_confirmed=pubkey_pre_confirmed)
+        PGPMixin.__init__(self)
 
     def _step_prepare(self):
         self.push('do_subscription')
         self.push('pubkey_confirmation')
         self.push('pubkey_checks')
+        self.push('pgp_prepare')
         self.push('confirmation_checks')
         self.push('verification_checks')
         self.push('sanity_checks')
@@ -107,7 +113,7 @@ class ConfirmSubscriptionPolicy(SubscriptionBase, VerificationMixin,
 @implementer(ISubscriptionWorkflow)
 class ModerationSubscriptionPolicy(SubscriptionBase, VerificationMixin,
                                    ModerationMixin, SetPubkeyMixin,
-                                   ConfirmPubkeyMixin):
+                                   ConfirmPubkeyMixin, PGPMixin):
     """"""
 
     name = 'pgp-policy-moderate'
@@ -133,12 +139,14 @@ class ModerationSubscriptionPolicy(SubscriptionBase, VerificationMixin,
         ModerationMixin.__init__(self, pre_approved=pre_approved)
         SetPubkeyMixin.__init__(self, pubkey=pubkey)
         ConfirmPubkeyMixin.__init__(self, pre_confirmed=pubkey_pre_confirmed)
+        PGPMixin.__init__(self)
 
     def _step_prepare(self):
         self.push('do_subscription')
         self.push('moderation_checks')
         self.push('pubkey_confirmation')
         self.push('pubkey_checks')
+        self.push('pgp_prepare')
         self.push('verification_checks')
         self.push('sanity_checks')
 
@@ -147,7 +155,8 @@ class ModerationSubscriptionPolicy(SubscriptionBase, VerificationMixin,
 @implementer(ISubscriptionWorkflow)
 class ConfirmModerationSubscriptionPolicy(SubscriptionBase, VerificationMixin,
                                           ConfirmationMixin, ModerationMixin,
-                                          SetPubkeyMixin, ConfirmPubkeyMixin):
+                                          SetPubkeyMixin, ConfirmPubkeyMixin,
+                                          PGPMixin):
     """"""
 
     name = 'pgp-policy-confirm-moderate'
@@ -175,12 +184,14 @@ class ConfirmModerationSubscriptionPolicy(SubscriptionBase, VerificationMixin,
         ModerationMixin.__init__(self, pre_approved=pre_approved)
         SetPubkeyMixin.__init__(self, pubkey=pubkey)
         ConfirmPubkeyMixin.__init__(self, pre_confirmed=pubkey_pre_confirmed)
+        PGPMixin.__init__(self)
 
     def _step_prepare(self):
         self.push('do_subscription')
         self.push('moderation_checks')
         self.push('pubkey_confirmation')
         self.push('pubkey_checks')
+        self.push('pgp_prepare')
         self.push('confirmation_checks')
         self.push('verification_checks')
         self.push('sanity_checks')
diff --git a/src/mailman_pgp/workflows/tests/test_base.py b/src/mailman_pgp/workflows/tests/test_base.py
index 913ab67..af00d01 100644
--- a/src/mailman_pgp/workflows/tests/test_base.py
+++ b/src/mailman_pgp/workflows/tests/test_base.py
@@ -38,7 +38,6 @@ from mailman_pgp.workflows.subscription import ConfirmSubscriptionPolicy
 
 
 class PubkeyMixinSetup():
-
     def setUp(self):
         with mm_transaction():
             self.mlist = create_list('test@example.com',
@@ -55,6 +54,27 @@ class PubkeyMixinSetup():
         self.sender = self.um.create_address('rsa-1024b@example.org')
 
 
+class TestPGPMixin(PubkeyMixinSetup, unittest.TestCase):
+    layer = PGPConfigLayer
+
+    def test_create_address(self):
+        workflow = ConfirmSubscriptionPolicy(self.mlist, self.sender,
+                                             pre_verified=True,
+                                             pre_confirmed=True)
+        workflow.run_thru('pgp_prepare')
+        pgp_address = PGPAddress.for_address(self.sender)
+        self.assertIsNotNone(pgp_address)
+
+    def test_address_existing(self):
+        workflow = ConfirmSubscriptionPolicy(self.mlist, self.sender,
+                                             pre_verified=True,
+                                             pre_confirmed=True)
+        with transaction() as t:
+            pgp_address = PGPAddress(self.sender)
+            t.add(pgp_address)
+        workflow.run_thru('pgp_prepare')
+
+
 class TestSetPubkeyMixin(PubkeyMixinSetup, unittest.TestCase):
     layer = PGPConfigLayer
 
@@ -75,60 +95,38 @@ class TestSetPubkeyMixin(PubkeyMixinSetup, unittest.TestCase):
                                              pre_verified=True,
                                              pre_confirmed=True)
         list(workflow)
-        with transaction() as t:
-            pgp_address = PGPAddress(self.sender)
+        with transaction():
+            pgp_address = PGPAddress.for_address(self.sender)
             pgp_address.key = self.sender_key.pubkey
-            t.add(pgp_address)
 
         receive_workflow = ConfirmSubscriptionPolicy(self.mlist)
         receive_workflow.token = workflow.token
         receive_workflow.restore()
         receive_workflow.run_thru('receive_key')
 
-        self.assertIsNotNone(receive_workflow.pubkey)
-        self.assertEqual(receive_workflow.pubkey.fingerprint,
-                         self.sender_key.pubkey.fingerprint)
-
-        receive_workflow.run_thru('pubkey_confirmation')
-        with patch.object(receive_workflow,
-                          '_step_send_key_confirm_request') as step:
-            next(receive_workflow)
-        step.assert_called_once_with()
-
-    def test_receive_key_no_address(self):
+    def test_set_pubkey(self):
         workflow = ConfirmSubscriptionPolicy(self.mlist, self.sender,
                                              pre_verified=True,
-                                             pre_confirmed=True)
-        list(workflow)
-        receive_workflow = ConfirmSubscriptionPolicy(self.mlist)
-        receive_workflow.token = workflow.token
-        receive_workflow.restore()
-        receive_workflow.run_thru('receive_key')
-
-        self.assertIsNone(receive_workflow.pubkey)
-        with patch.object(receive_workflow,
-                          '_step_send_key_request') as step:
-            next(receive_workflow)
-        step.assert_called_once_with()
+                                             pre_confirmed=True,
+                                             pubkey=self.sender_key.pubkey)
+        workflow.run_thru('pubkey_checks')
+        pgp_address = PGPAddress.for_address(self.sender)
+        self.assertIsNotNone(pgp_address)
+        self.assertIsNotNone(pgp_address.key)
+        self.assertEqual(pgp_address.key_fingerprint,
+                         self.sender_key.fingerprint)
 
-    def test_receive_key_pubkey_confirmed(self):
+    def test_pubkey_set(self):
         workflow = ConfirmSubscriptionPolicy(self.mlist, self.sender,
                                              pre_verified=True,
-                                             pre_confirmed=True,
-                                             pubkey_pre_confirmed=True)
-        list(workflow)
+                                             pre_confirmed=True)
         with transaction() as t:
             pgp_address = PGPAddress(self.sender)
             pgp_address.key = self.sender_key.pubkey
             t.add(pgp_address)
-
-        receive_workflow = ConfirmSubscriptionPolicy(self.mlist)
-        receive_workflow.token = workflow.token
-        receive_workflow.restore()
-        receive_workflow.run_thru('pubkey_confirmation')
-        with patch.object(receive_workflow, '_step_do_subscription') as step:
-            next(receive_workflow)
-        step.assert_called_once_with()
+        workflow.run_thru('pubkey_checks')
+        self.assertEqual(pgp_address.key_fingerprint,
+                         self.sender_key.fingerprint)
 
 
 class TestConfirmPubkeyMixin(PubkeyMixinSetup, unittest.TestCase):
@@ -168,11 +166,6 @@ class TestConfirmPubkeyMixin(PubkeyMixinSetup, unittest.TestCase):
                                              pubkey_pre_confirmed=False)
         list(workflow)
 
-        with transaction() as t:
-            pgp_address = PGPAddress(self.sender)
-            pgp_address.key = self.sender_key.pubkey
-            t.add(pgp_address)
-
         receive_workflow = ConfirmSubscriptionPolicy(self.mlist)
         receive_workflow.token = workflow.token
         receive_workflow.restore()