crypto: Use session-key to decrypt messages if present in the index

notmuch caches the OpenPGP session keys if configured to do so. See
index.decrypt on:
https://notmuchmail.org/manpages/notmuch-config-1/

Using the cached session key decryption of messages can be done without
the need of having the private OpenPGP key. There is some speed up on
decryption, mostly notable on long encrypted threads.

5 files changed, 60 insertions, 15 deletions

diff --git a/alot/crypto.py b/alot/crypto.py
index a05db2fc..38b8727a 100644
--- a/alot/crypto.py
+++ b/alot/crypto.py
@@ -206,15 +206,22 @@ def verify_detached(message, signature):
         raise GPGProblem(str(e), code=e.getcode())
 
 
-def decrypt_verify(encrypted):
+def decrypt_verify(encrypted, session_keys=None):
     """Decrypts the given ciphertext string and returns both the
     signatures (if any) and the plaintext.
 
     :param bytes encrypted: the mail to decrypt
+    :param list[str] session_keys: a list OpenPGP session keys
     :returns: the signatures and decrypted plaintext data
     :rtype: tuple[list[gpg.resuit.Signature], str]
     :raises: :class:`~alot.errors.GPGProblem` if the decryption fails
     """
+    if session_keys is not None:
+        try:
+            return _decrypt_verify_session_keys(encrypted, session_keys)
+        except GPGProblem:
+            pass
+
     ctx = gpg.core.Context()
     try:
         plaintext, _, verify_result = ctx.decrypt(encrypted, verify=True)
@@ -228,6 +235,28 @@ def decrypt_verify(encrypted):
     return sigs, plaintext
 
 
+def _decrypt_verify_session_keys(encrypted, session_keys):
+    """Decrypts the given ciphertext string using the session_keys
+    and returns both the signatures (if any) and the plaintext.
+
+    :param bytes encrypted: the mail to decrypt
+    :param list[str] session_keys: a list OpenPGP session keys
+    :returns: the signatures and decrypted plaintext data
+    :rtype: tuple[list[gpg.resuit.Signature], str]
+    :raises: :class:`~alot.errors.GPGProblem` if the decryption fails
+    """
+    for key in session_keys:
+        ctx = gpg.core.Context()
+        ctx.set_ctx_flag("override-session-key", key)
+        try:
+            (plaintext, _, verify_result) = ctx.decrypt(
+                    encrypted, verify=True)
+        except gpg.errors.GPGMEError as e:
+            continue
+        return verify_result.signatures, plaintext
+    raise GPGProblem("No valid session key", code=GPGCode.NOT_FOUND)
+
+
 def validate_key(key, sign=False, encrypt=False):
     """Assert that a key is valide and optionally that it can be used for
     signing or encrypting.  Raise GPGProblem otherwise.
diff --git a/alot/db/message.py b/alot/db/message.py
index 3c860cac..ae06acc6 100644
--- a/alot/db/message.py
+++ b/alot/db/message.py
@@ -54,6 +54,11 @@ class Message(object):
         self._attachments = None  # will be read upon first use
         self._tags = set(msg.get_tags())
 
+        self._session_keys = []
+        for name, value in msg.get_properties("session-key", exact=True):
+            if name == "session-key":
+                self._session_keys.append(value)
+
         try:
             sender = decode_header(msg.get_header('From'))
             if not sender:
@@ -102,7 +107,8 @@ class Message(object):
         if not self._email:
             try:
                 with open(path, 'rb') as f:
-                    self._email = utils.decrypted_message_from_bytes(f.read())
+                    self._email = utils.decrypted_message_from_bytes(
+                            f.read(), self._session_keys)
             except IOError:
                 self._email = email.message_from_string(
                     warning, policy=email.policy.SMTP)
diff --git a/alot/db/utils.py b/alot/db/utils.py
index c883bb42..abbdd4ea 100644
--- a/alot/db/utils.py
+++ b/alot/db/utils.py
@@ -140,7 +140,7 @@ def _handle_signatures(original, message, params):
     add_signature_headers(original, sigs, malformed)
 
 
-def _handle_encrypted(original, message):
+def _handle_encrypted(original, message, session_keys=None):
     """Handle encrypted messages helper.
 
     RFC 3156 is quite strict:
@@ -155,6 +155,8 @@ def _handle_encrypted(original, message):
     :type original: :class:`email.message.Message`
     :param message: The multipart/signed payload to verify
     :type message: :class:`email.message.Message`
+    :param session_keys: a list OpenPGP session keys
+    :type session_keys: [str]
     """
     malformed = False
 
@@ -172,14 +174,14 @@ def _handle_encrypted(original, message):
         # This should be safe because PGP uses US-ASCII characters only
         payload = message.get_payload(1).get_payload().encode('ascii')
         try:
-            sigs, d = crypto.decrypt_verify(payload)
+            sigs, d = crypto.decrypt_verify(payload, session_keys)
         except GPGProblem as e:
             # signature verification failures end up here too if the combined
             # method is used, currently this prevents the interpretation of the
             # recovered plain text mail. maybe that's a feature.
             malformed = str(e)
         else:
-            n = decrypted_message_from_bytes(d)
+            n = decrypted_message_from_bytes(d, session_keys)
 
             # add the decrypted message to message. note that n contains all
             # the attachments, no need to walk over n here.
@@ -214,7 +216,7 @@ def _handle_encrypted(original, message):
         original.attach(content)
 
 
-def decrypted_message_from_file(handle):
+def decrypted_message_from_file(handle, session_keys=None):
     '''Reads a mail from the given file-like object and returns an email
     object, very much like email.message_from_file. In addition to
     that OpenPGP encrypted data is detected and decrypted. If this
@@ -222,18 +224,21 @@ def decrypted_message_from_file(handle):
     message are added to the returned message object.
 
     :param handle: a file-like object
+    :param session_keys: a list OpenPGP session keys
     :returns: :class:`email.message.Message` possibly augmented with
               decrypted data
     '''
-    return decrypted_message_from_message(email.message_from_file(handle))
+    return decrypted_message_from_message(email.message_from_file(handle),
+                                          session_keys)
 
 
-def decrypted_message_from_message(m):
+def decrypted_message_from_message(m, session_keys=None):
     '''Detect and decrypt OpenPGP encrypted data in an email object. If this
     succeeds, any mime messages found in the recovered plaintext
     message are added to the returned message object.
 
     :param m: an email object
+    :param session_keys: a list OpenPGP session keys
     :returns: :class:`email.message.Message` possibly augmented with
               decrypted data
     '''
@@ -253,7 +258,7 @@ def decrypted_message_from_message(m):
         elif (m.get_content_subtype() == 'encrypted' and
               p.get('protocol') == _APP_PGP_ENC and
               'Version: 1' in m.get_payload(0).get_payload()):
-            _handle_encrypted(m, m)
+            _handle_encrypted(m, m, session_keys)
 
         # It is also possible to put either of the abov into a multipart/mixed
         # segment
@@ -268,12 +273,12 @@ def decrypted_message_from_message(m):
                     _handle_signatures(m, sub, p)
                 elif (sub.get_content_subtype() == 'encrypted' and
                       p.get('protocol') == _APP_PGP_ENC):
-                    _handle_encrypted(m, sub)
+                    _handle_encrypted(m, sub, session_keys)
 
     return m
 
 
-def decrypted_message_from_string(s):
+def decrypted_message_from_string(s, session_keys=None):
     '''Reads a mail from the given string. This is the equivalent of
     :func:`email.message_from_string` which does nothing but to wrap
     the given string in a StringIO object and to call
@@ -283,16 +288,18 @@ def decrypted_message_from_string(s):
     details.
 
     '''
-    return decrypted_message_from_file(io.StringIO(s))
+    return decrypted_message_from_file(io.StringIO(s), session_keys)
 
 
-def decrypted_message_from_bytes(bytestring):
+def decrypted_message_from_bytes(bytestring, session_keys=None):
     """Create a Message from bytes.
 
     :param bytes bytestring: an email message as raw bytes
+    :param session_keys: a list OpenPGP session keys
     """
     return decrypted_message_from_message(
-        email.message_from_bytes(bytestring, policy=email.policy.SMTP))
+        email.message_from_bytes(bytestring, policy=email.policy.SMTP),
+        session_keys)
 
 
 def extract_headers(mail, headers=None):
diff --git a/setup.py b/setup.py
index 6c322404..beb17546 100755
--- a/setup.py
+++ b/setup.py
@@ -45,7 +45,7 @@ setup(
             ['alot = alot.__main__:main'],
     },
     install_requires=[
-        'notmuch>=0.26',
+        'notmuch>=0.27',
         'urwid>=1.3.0',
         'urwidtrees>=1.0',
         'twisted>=10.2.0',
diff --git a/tests/db/message_test.py b/tests/db/message_test.py
index caa70a44..29ed5ee6 100644
--- a/tests/db/message_test.py
+++ b/tests/db/message_test.py
@@ -57,6 +57,9 @@ class MockNotmuchMessage(object):
     def get_tags(self):
         return self.mock_tags
 
+    def get_properties(self, prop, exact=False):
+        return []
+
 
 class TestMessage(unittest.TestCase):