plugin_e2ee: allow decrypt call to be async

Signed-off-by: Maxime “pep” Buquet <pep@bouah.net>

1 files changed, 35 insertions, 6 deletions

diff --git a/poezio/plugin_e2ee.py b/poezio/plugin_e2ee.py
index 183350ea..1b393841 100644
--- a/poezio/plugin_e2ee.py
+++ b/poezio/plugin_e2ee.py
@@ -148,9 +148,9 @@ class E2EEPlugin(BasePlugin):
 
         # Ensure decryption is done before everything, so that other handlers
         # don't have to know about the encryption mechanism.
-        self.api.add_event_handler('muc_msg', self._decrypt, priority=0)
-        self.api.add_event_handler('conversation_msg', self._decrypt, priority=0)
-        self.api.add_event_handler('private_msg', self._decrypt, priority=0)
+        self.api.add_event_handler('muc_msg', self._decrypt_wrapper, priority=0)
+        self.api.add_event_handler('conversation_msg', self._decrypt_wrapper, priority=0)
+        self.api.add_event_handler('private_msg', self._decrypt_wrapper, priority=0)
 
         # Ensure encryption is done after everything, so that whatever can be
         # encrypted is encrypted, and no plain element slips in.
@@ -345,6 +345,28 @@ class E2EEPlugin(BasePlugin):
         except NothingToEncrypt:
             return stanza
         except Exception as exc:
+            jid = stanza['from']
+            tab = self.core.tabs.by_name_and_class(jid, ChatTab)
+            msg = ' \n\x19%s}Could not decrypt message: %s' % (
+                dump_tuple(get_theme().COLOR_CHAR_NACK),
+                exc,
+            )
+            # XXX: check before commit. Do we not nack in MUCs?
+            if tab and not isinstance(tab, MucTab):
+                tab.nack_message(msg, stanza['id'], stanza['to'])
+            # TODO: display exceptions to the user properly
+            log.error('Exception in encrypt:', exc_info=True)
+            return None
+        return result
+
+    async def _decrypt_wrapper(self, stanza: StanzaBase, tab: ChatTabs) -> Optional[StanzaBase]:
+        """
+        Wrapper around _decrypt() to handle errors and display the message after encryption.
+        """
+        try:
+            # pylint: disable=unexpected-keyword-arg
+            result = await self._decrypt(stanza, tab, passthrough=True)
+        except Exception as exc:
             jid = stanza['to']
             tab = self.core.tabs.by_name_and_class(jid, ChatTab)
             msg = ' \n\x19%s}Could not send message: %s' % (
@@ -355,11 +377,11 @@ class E2EEPlugin(BasePlugin):
             if tab and not isinstance(tab, MucTab):
                 tab.nack_message(msg, stanza['id'], stanza['from'])
             # TODO: display exceptions to the user properly
-            log.error('Exception in encrypt:', exc_info=True)
+            log.error('Exception in decrypt:', exc_info=True)
             return None
         return result
 
-    def _decrypt(self, message: Message, tab: ChatTabs) -> None:
+    async def _decrypt(self, message: Message, tab: ChatTabs, passthrough: bool = True) -> None:
 
         has_eme = False
         if message.xml.find('{%s}%s' % (EME_NS, EME_TAG)) is not None and \
@@ -397,7 +419,14 @@ class E2EEPlugin(BasePlugin):
             if user is not None:
                 jid = user.jid or None
 
-        self.decrypt(message, jid, tab)
+        # Call the enabled encrypt method
+        func = self.decrypt
+        if iscoroutinefunction(func):
+            # pylint: disable=unexpected-keyword-arg
+            await func(message, jid, tab, passthrough=True)
+        else:
+            # pylint: disable=unexpected-keyword-arg
+            func(message, jid, tab)
 
         log.debug('Decrypted %s message: %r', self.encryption_name, message['body'])
         return None