Add repair to pixelated maintenance

- Issue #468
- Call with: pixelated-maintenance repair

5 files changed, 69 insertions, 3 deletions

diff --git a/service/pixelated/adapter/mailstore/maintenance/__init__.py b/service/pixelated/adapter/mailstore/maintenance/__init__.py
index eaac5814..edc442c2 100644
--- a/service/pixelated/adapter/mailstore/maintenance/__init__.py
+++ b/service/pixelated/adapter/mailstore/maintenance/__init__.py
@@ -14,6 +14,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with Pixelated. If not, see <http://www.gnu.org/licenses/>.
 from leap.keymanager.keys import KEY_TYPE_KEY, KEY_PRIVATE_KEY, KEY_ID_KEY, KEY_ADDRESS_KEY
+from leap.keymanager.openpgp import OpenPGPKey
 
 from twisted.internet import defer
 import logging
@@ -35,6 +36,14 @@ def _is_private_key_doc(doc):
     return _is_key_doc(doc) and doc.content.get(KEY_PRIVATE_KEY, False)
 
 
+def _is_active_key_doc(doc):
+    return _is_key_doc(doc) and doc.content.get(KEY_TYPE_KEY, None) == TYPE_OPENPGP_ACTIVE
+
+
+def _is_public_key(doc):
+    return _is_key_doc(doc) and not doc.content.get(KEY_PRIVATE_KEY, False)
+
+
 def _key_id(doc):
     return doc.content.get(KEY_ID_KEY, None)
 
@@ -58,5 +67,33 @@ class SoledadMaintenance(object):
                 logger.warn('Deleting doc %s for key %s of <%s>' % (doc.doc_id, _key_id(doc), _address(doc)))
                 yield self._soledad.delete_doc(doc)
 
+        yield self._repair_missing_active_docs(docs, private_key_ids)
+
+    @defer.inlineCallbacks
+    def _repair_missing_active_docs(self, docs, private_key_ids):
+        missing = self._missing_active_docs(docs, private_key_ids)
+        for key_id in missing:
+            emails = self._emails_for_key_id(docs, key_id)
+            for email in emails:
+                logger.warn('Re-creating active doc for key %s, email %s' % (key_id, email))
+                yield self._soledad.create_doc_from_json(OpenPGPKey(email, key_id=key_id, private=False).get_active_json(email))
+
     def _key_ids_with_private_key(self, docs):
         return [doc.content[KEY_ID_KEY] for doc in docs if _is_private_key_doc(doc)]
+
+    def _missing_active_docs(self, docs, private_key_ids):
+        active_doc_ids = self._active_docs_for_key_id(docs)
+
+        return set([private_key_id for private_key_id in private_key_ids if private_key_id not in active_doc_ids])
+
+    def _emails_for_key_id(self, docs, key_id):
+        for doc in docs:
+            if _is_private_key_doc(doc) and _key_id(doc) == key_id:
+                email = _address(doc)
+                if isinstance(email, list):
+                    return email
+                else:
+                    return [email]
+
+    def _active_docs_for_key_id(self, docs):
+        return [doc.content[KEY_ID_KEY] for doc in docs if _is_active_key_doc(doc) and _is_public_key(doc)]
diff --git a/service/pixelated/config/arguments.py b/service/pixelated/config/arguments.py
index fa7fdae4..7a7abe49 100644
--- a/service/pixelated/config/arguments.py
+++ b/service/pixelated/config/arguments.py
@@ -43,6 +43,7 @@ def parse_maintenance_args():
 
     subparsers.add_parser('dump-soledad', help='dump the soledad database')
     subparsers.add_parser('sync', help='sync the soledad database')
+    subparsers.add_parser('repair', help='repair database if possible')
 
     return parser.parse_args()
 
diff --git a/service/pixelated/config/leap.py b/service/pixelated/config/leap.py
index c1280756..4dcb18f4 100644
--- a/service/pixelated/config/leap.py
+++ b/service/pixelated/config/leap.py
@@ -13,7 +13,8 @@ def initialize_leap(leap_provider_cert,
                     leap_provider_cert_fingerprint,
                     credentials_file,
                     organization_mode,
-                    leap_home):
+                    leap_home,
+                    initial_sync=True):
     init_monkeypatches()
     events_server.ensure_server()
     provider, username, password = credentials.read(organization_mode, credentials_file)
@@ -24,7 +25,8 @@ def initialize_leap(leap_provider_cert,
     LeapCertificate(provider).setup_ca_bundle()
     leap_session = LeapSessionFactory(provider).create(username, password)
 
-    leap_session = yield leap_session.initial_sync()
+    if initial_sync:
+        leap_session = yield leap_session.initial_sync()
 
     defer.returnValue(leap_session)
 
diff --git a/service/pixelated/maintenance.py b/service/pixelated/maintenance.py
index ae320049..41e749d7 100644
--- a/service/pixelated/maintenance.py
+++ b/service/pixelated/maintenance.py
@@ -18,6 +18,7 @@ import logging
 from mailbox import Maildir
 from twisted.internet import reactor, defer
 from twisted.internet.threads import deferToThread
+from pixelated.adapter.mailstore.maintenance import SoledadMaintenance
 from pixelated.config.leap import initialize_leap
 from pixelated.config import logger, arguments
 
@@ -36,7 +37,8 @@ def initialize():
             args.leap_provider_cert_fingerprint,
             args.credentials_file,
             organization_mode=False,
-            leap_home=args.leap_home)
+            leap_home=args.leap_home,
+            initial_sync=False)
 
         execute_command(args, leap_session)
 
@@ -91,6 +93,9 @@ def add_command_callback(args, prepareDeferred, finalizeDeferred):
     elif args.command == 'sync':
         # nothing to do here, sync is already part of the chain
         prepareDeferred.chainDeferred(finalizeDeferred)
+    elif args.command == 'repair':
+        prepareDeferred.addCallback(repair)
+        prepareDeferred.chainDeferred(finalizeDeferred)
     else:
         print 'Unsupported command: %s' % args.command
         prepareDeferred.chainDeferred(finalizeDeferred)
@@ -176,6 +181,15 @@ def dump_soledad(args):
     defer.returnValue(args)
 
 
+@defer.inlineCallbacks
+def repair(args):
+    leap_session, soledad = args
+
+    yield SoledadMaintenance(soledad).repair()
+
+    defer.returnValue(args)
+
+
 def shutdown(args):
     # time.sleep(30)
     reactor.stop()
diff --git a/service/test/unit/adapter/mailstore/maintenance/test_soledad_maintenance.py b/service/test/unit/adapter/mailstore/maintenance/test_soledad_maintenance.py
index dc5d9d6c..9a89d62b 100644
--- a/service/test/unit/adapter/mailstore/maintenance/test_soledad_maintenance.py
+++ b/service/test/unit/adapter/mailstore/maintenance/test_soledad_maintenance.py
@@ -87,6 +87,18 @@ class TestSoledadMaintenance(unittest.TestCase):
 
         verify(soledad, never).delete_doc(other_doc)
 
+    @defer.inlineCallbacks
+    def test_repair_recreates_public_key_active_doc_if_necessary(self):
+        soledad = mock()
+
+        private_key = self._private_key(SOME_EMAIL_ADDRESS, SOME_KEY_ID)
+        private_key_doc = SoledadDocument(doc_id='some_doc', json=private_key.get_json())
+        when(soledad).get_all_docs().thenReturn(defer.succeed((1, [private_key_doc])))
+
+        yield SoledadMaintenance(soledad).repair()
+
+        verify(soledad).create_doc_from_json('{"key_id": "4914254E384E264C", "tags": ["keymanager-active"], "type": "OpenPGPKey-active", "private": false, "address": "foo@example.tld"}')
+
     def _public_key(self, address, keyid):
         return self._gpgkey(address, keyid, private=False)