4 files changed, 91 insertions, 71 deletions

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()