echo_bot: Also match on heartbeats

The default slixmpp 'message' handler only matches on messages including
a body, which causes issues with heartbeats as they should be processed
by the OMEMO lib as well.

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

examples/echo_client.py

@@ -21,6 +21,8 @@ from argparse import ArgumentParser
 from slixmpp import ClientXMPP, JID
 from slixmpp.exceptions import IqTimeout, IqError
 from slixmpp.stanza import Message
+from slixmpp.xmlstream.handler import Callback
+from slixmpp.xmlstream.matcher import MatchXPath
 import slixmpp_omemo
 from slixmpp_omemo import PluginCouldNotLoad, MissingOwnKey, EncryptionPrepareException
 from slixmpp_omemo import UndecidedException, UntrustedException, NoAvailableSession
@@ -54,7 +56,11 @@ class EchoBot(ClientXMPP):
         self.cmd_re: re.Pattern = re.compile('^%s(?P<command>\w+)(?:\s+(?P<args>.*))?' % self.cmd_prefix)
 
         self.add_event_handler("session_start", self.start)
-        self.add_event_handler("message", self.message_handler)
+        self.add_event_handler("message", self.plain_handler)  # Matches any message with body
+        self.register_handler(Callback('Encrypted',  # Matches any encrypted message
+            MatchXPath(f'{{{self.default_ns}}}message/{{{self.eme_ns}}}encrypted'),
+            self.encrypted_handler,
+        ))
 
     def start(self, _event) -> None:
         """
@@ -112,10 +118,27 @@ class EchoBot(ClientXMPP):
         body = '''Debug level set to 'error'.'''
         return await self.encrypted_reply(mto, mtype, body)
 
-    def message_handler(self, msg: Message) -> None:
+    def plain_handler(self, msg: Message) -> None:
-        asyncio.ensure_future(self.message(msg))
+        asyncio.ensure_future(self.plain_message(msg))
 
-    async def message(self, msg: Message, allow_untrusted: bool = False) -> None:
+    async def plain_message(self, msg: Message) -> None:
+        if self['xep_0384'].is_encrypted(msg):  # Already handled by 'encrypted_handler'
+            return None
+
+        mfrom = mto = msg['from']
+        mtype = msg['type']
+
+        if mtype not in ('chat', 'normal'):
+            return None
+
+        if self.debug_level == LEVEL_DEBUG:
+            await self.plain_reply(mto, mtype, 'Echo unencrypted message:%(body)s' % msg)
+        return None
+
+    def encrypted_handler(self, msg: Message) -> None:
+        asyncio.ensure_future(self.encrypted_message(msg))
+
+    async def encrypted_message(self, msg: Message, allow_untrusted: bool = False) -> None:
         """
         Process incoming message stanzas. Be aware that this also
         includes MUC messages and error messages. It is usually
@@ -133,11 +156,6 @@ class EchoBot(ClientXMPP):
         if mtype not in ('chat', 'normal'):
             return None
 
-        if not self['xep_0384'].is_encrypted(msg):
-            if self.debug_level == LEVEL_DEBUG:
-                await self.plain_reply(mto, mtype, 'Echo unencrypted message:%(body)s' % msg)
-            return None
-
         try:
             encrypted = msg['omemo_encrypted']
             body = await self['xep_0384'].decrypt_message(encrypted, mfrom, allow_untrusted)