Echo_bot: Use expect_problems and handle MissingBundleException

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

examples/echo_client.py

@@ -18,10 +18,12 @@ from getpass import getpass
 from argparse import ArgumentParser
 
 from slixmpp import ClientXMPP, JID
+from slixmpp.exceptions import IqTimeout, IqError
 from slixmpp.stanza import Message
+import slixmpp_omemo
 from slixmpp_omemo import PluginCouldNotLoad, MissingOwnKey, EncryptionPrepareException
 from slixmpp_omemo import UndecidedException, UntrustedException, NoAvailableSession
-import slixmpp_omemo
+from omemo.exceptions import MissingBundleException
 
 log = logging.getLogger(__name__)
 
@@ -158,6 +160,8 @@ class EchoBot(ClientXMPP):
         msg['eme']['namespace'] = self.eme_ns
         msg['eme']['name'] = self['xep_0380'].mechanisms[self.eme_ns]
 
+        expect_problems = {}  # type: Optional[Dict[JID, List[int]]]
+
         while True:
             try:
                 # `encrypt_message` excepts the plaintext to be sent, a list of
@@ -171,7 +175,7 @@ class EchoBot(ClientXMPP):
                 #
                 # TODO: Document expect_problems
                 recipients = [mto]
-                encrypt = await self['xep_0384'].encrypt_message(body, recipients)
+                encrypt = await self['xep_0384'].encrypt_message(body, recipients, expect_problems)
                 msg.append(encrypt)
                 return msg.send()
             except UndecidedException as exn:
@@ -180,7 +184,25 @@ class EchoBot(ClientXMPP):
                 # This is where you prompt your user to ask what to do. In
                 # this bot we will automatically trust undecided recipients.
                 self['xep_0384'].trust(exn.bare_jid, exn.device, exn.ik)
-            # TODO: catch NoEligibleDevicesException and MissingBundleException
+            # TODO: catch NoEligibleDevicesException
+            except EncryptionPrepareException as exn:
+                # TODO: We might need to bail out here if errors are the same?
+                for error in exn.errors:
+                    if isinstance(error, MissingBundleException):
+                        self.plain_reply(
+                            original_msg,
+                            'Could not find keys for device "%d" of recipient "%s". Skipping.' %
+                            (error.device, error.bare_jid),
+                        )
+                        jid = JID(error.bare_jid)
+                        device_list = expect_problems.setdefault(jid, [])
+                        device_list.append(error.device)
+            except (IqError, IqTimeout) as exn:
+                self.plain_reply(
+                    original_msg,
+                    'An error occured while fetching information on a recipient.\n%r' % exn,
+                )
+                return None
             except Exception as exn:
                 await self.plain_reply(
                     original_msg,
@@ -190,6 +212,7 @@ class EchoBot(ClientXMPP):
 
         return None
 
+
 if __name__ == '__main__':
     # Setup the command line arguments.
     parser = ArgumentParser(description=EchoBot.__doc__)