# HG changeset patch
# User Goffi <goffi@goffi.org>
# Date 1534004695 -7200
# Node ID 0bef44f8e8ca293c500b075b62c4516dd8b28091
# Parent  661f66d412158dd480efd74a2eb96c2c5e0da460
plugin XEP-0384: PEP handling + mark as encrypted:

devices are gotten using PEP and +notify, and bundles are then retrieved, and kept in cache.
This allow to remove the test code which was always requesting devices and bundles before
sending a message. For now OMEMO plugin will only work with people in roster, this will be
fixed in the future.

diff -r 661f66d41215 -r 0bef44f8e8ca sat/plugins/plugin_xep_0384.py
--- a/sat/plugins/plugin_xep_0384.py	Sat Aug 11 18:24:55 2018 +0200
+++ b/sat/plugins/plugin_xep_0384.py	Sat Aug 11 18:24:55 2018 +0200
@@ -47,7 +47,7 @@
     C.PI_IMPORT_NAME: u"OMEMO",
     C.PI_TYPE: u"SEC",
     C.PI_PROTOCOLS: [u"XEP-0384"],
-    C.PI_DEPENDENCIES: [u"XEP-0280", u"XEP-0334", u"XEP-0060"],
+    C.PI_DEPENDENCIES: [u"XEP-0163", u"XEP-0280", u"XEP-0334", u"XEP-0060"],
     C.PI_MAIN: u"OMEMO",
     C.PI_HANDLER: u"no",
     C.PI_DESCRIPTION: _(u"""Implementation of OMEMO"""),
@@ -266,6 +266,8 @@
         host.trigger.add("MessageReceived", self._messageReceivedTrigger, priority=100050)
         host.trigger.add("sendMessageData", self._sendMessageDataTrigger)
         self.host.registerEncryptionPlugin(self, u"OMEMO", NS_OMEMO, 100)
+        pep = host.plugins['XEP-0163']
+        pep.addPEPEvent("OMEMO_DEVICES", NS_OMEMO_DEVICES, self.onNewDevices)
 
     @defer.inlineCallbacks
     def profileConnected(self, client):
@@ -290,6 +292,7 @@
 
         omemo_storage = OmemoStorage(client, device_id, persistent_dict)
         omemo_session = yield OmemoSession.create(client, omemo_storage, device_id)
+        client._xep_0384_cache = {}
         client._xep_0384_session = omemo_session
         client._xep_0384_device_id = device_id
         yield omemo_session.newDeviceList(devices, client.jid)
@@ -303,24 +306,13 @@
 
     # devices
 
-    @defer.inlineCallbacks
-    def getDevices(self, client, entity_jid=None):
-        """Retrieve list of registered OMEMO devices
+    def parseDevices(self, items):
+        """Parse devices found in items
 
-        @param entity_jid(jid.JID, None): get devices from this entity
-            None to get our own devices
-        @return (set(int)): list of devices
+        @param items(iterable[domish.Element]): items as retrieved by getItems
+        @return set[int]: parsed devices
         """
-        if entity_jid is not None:
-            assert not entity_jid.resource
         devices = set()
-        try:
-            items, metadata = yield self._p.getItems(client, entity_jid, NS_OMEMO_DEVICES)
-        except error.StanzaError as e:
-            if e.condition == 'item-not-found':
-                log.info(_(u"there is no node to handle OMEMO devices"))
-                defer.returnValue(devices)
-
         if len(items) > 1:
             log.warning(_(u"OMEMO devices list is stored in more that one items, "
                           u"this is not expected"))
@@ -341,6 +333,26 @@
                         device_id=device_elt['id']))
                 else:
                     devices.add(device_id)
+        return devices
+
+    @defer.inlineCallbacks
+    def getDevices(self, client, entity_jid=None):
+        """Retrieve list of registered OMEMO devices
+
+        @param entity_jid(jid.JID, None): get devices from this entity
+            None to get our own devices
+        @return (set(int)): list of devices
+        """
+        if entity_jid is not None:
+            assert not entity_jid.resource
+        try:
+            items, metadata = yield self._p.getItems(client, entity_jid, NS_OMEMO_DEVICES)
+        except error.StanzaError as e:
+            if e.condition == 'item-not-found':
+                log.info(_(u"there is no node to handle OMEMO devices"))
+                defer.returnValue(set())
+
+        devices = self.parseDevices(items)
         defer.returnValue(devices)
 
     def setDevicesEb(self, failure_):
@@ -466,14 +478,33 @@
         d.addErrback(self.setBundleEb)
         return d
 
+    ## PEP node events callbacks
+
+    @defer.inlineCallbacks
+    def onNewDevices(self, itemsEvent, profile):
+        client = self.host.getClient(profile)
+        cache = client._xep_0384_cache
+        omemo_session = client._xep_0384_session
+        entity = itemsEvent.sender
+        entity_cache = cache.setdefault(entity, {})
+        devices = self.parseDevices(itemsEvent.items)
+        omemo_session.newDeviceList(devices, entity)
+        missing_devices = devices.difference(entity_cache.keys())
+        if missing_devices:
+            missing_bundles = yield self.getBundles(client, entity, missing_devices)
+            entity_cache.update(missing_bundles)
+
     ## triggers
 
     @defer.inlineCallbacks
     def encryptMessage(self, client, entity_bare_jid, message):
         omemo_session = client._xep_0384_session
-        devices = yield self.getDevices(client, entity_bare_jid)
-        omemo_session.newDeviceList(devices, entity_bare_jid)
-        bundles = yield self.getBundles(client, entity_bare_jid, devices)
+        try:
+            bundles = client._xep_0384_cache[entity_bare_jid]
+        except KeyError:
+            raise exceptions.NotFound(_(u"No OMEMO encryption information found for this"
+                                        u"contact ({entity})".format(
+                                        entity=entity_bare_jid.full())))
         encrypted = yield omemo_session.encryptMessage(
             entity_bare_jid,
             message,
@@ -547,6 +578,7 @@
         message_elt.children.remove(encrypted_elt)
         if plaintext:
             message_elt.addElement("body", content=plaintext.decode('utf-8'))
+        post_treat.addCallback(client.encryption.markAsEncrypted)
         defer.returnValue(True)
 
     @defer.inlineCallbacks