plugin_e2ee: let already encrypted messages through

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

1 files changed, 27 insertions, 13 deletions

diff --git a/poezio/plugin_e2ee.py b/poezio/plugin_e2ee.py
index 02c19f4c..178a3b51 100644
--- a/poezio/plugin_e2ee.py
+++ b/poezio/plugin_e2ee.py
@@ -448,9 +448,19 @@ class E2EEPlugin(BasePlugin):
         if tab is None:  # Possible message sent directly by the e2ee lib?
             log.debug(
                 'A message we do not have a tab for '
-                'is being sent to \'%s\'. Abort.', message['to'],
+                'is being sent to \'%s\'. \n%r.', message['to'], message,
             )
-            return None
+
+        # Is this message already encrypted? Do we need to do all these
+        # checks? Possibly an OMEMO heartbeat.
+        has_encrypted_tag = False
+        if self.encrypted_tags is not None:
+            for (namespace, tag) in self.encrypted_tags:
+                if message.xml.find('{%s}%s' % (namespace, tag)) is not None:
+                    # TODO: count all encrypted tags.
+                    has_encrypted_tag = True
+                    log.debug('Message already contains an encrypted tag')
+                    break
 
         parent = None
         if isinstance(tab, PrivateTab):
@@ -492,22 +502,26 @@ class E2EEPlugin(BasePlugin):
 
         # Drop all messages that don't contain a body if the plugin doesn't do
         # Stanza Encryption
-        if not self.stanza_encryption and not has_body:
+        if not self.stanza_encryption and not has_encrypted_tag and not has_body:
             log.debug(
-                '%s plugin: Dropping message as it contains no body, and '
-                'not doesn\'t do stanza encryption',
+                '%s plugin: Dropping message as it\'s not already encrypted, '
+                'contains no body, and doesn\'t do stanza encryption',
                 self.encryption_name,
             )
             return None
 
-        # Call the enabled encrypt method
-        func = self._enabled_tabs[tab.jid]
-        if iscoroutinefunction(func):
-            # pylint: disable=unexpected-keyword-arg
-            await func(message, jids, tab, passthrough=True)
-        else:
-            # pylint: disable=unexpected-keyword-arg
-            func(message, jids, tab, passthrough=True)
+        if tab and not has_encrypted_tag:
+            if not self._encryption_enabled(tab.jid):
+                raise NothingToEncrypt()
+
+            # Call the enabled encrypt method
+            func = self._enabled_tabs[tab.jid]
+            if iscoroutinefunction(func):
+                # pylint: disable=unexpected-keyword-arg
+                await func(message, jids, tab, passthrough=True)
+            else:
+                # pylint: disable=unexpected-keyword-arg
+                func(message, jids, tab, passthrough=True)
 
         if has_body:
             # Only add EME tag if the message has a body.