1 files changed, 29 insertions, 19 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