From 3322d17aa4ed0b00b19ba2978410fe2e0c6d88e5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maxime=20=E2=80=9Cpep=E2=80=9D=20Buquet?= <pep@bouah.net>
Date: Mon, 14 May 2018 12:16:14 +0200
Subject: [PATCH] xep_0384: Handle device list

- Split get/store and common operations.
- Always try to set at init time
- Add ourselves in the list if not present
- Correctly cast from/to str/int
---
 plugin.py | 55 +++++++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 45 insertions(+), 10 deletions(-)

diff --git a/plugin.py b/plugin.py
index ed2bb36..52094cb 100644
--- a/plugin.py
+++ b/plugin.py
@@ -9,8 +9,10 @@
 import logging
 
 import base64
+import asyncio
 from slixmpp.plugins.xep_0384.stanza import OMEMO_BASE_NS
 from slixmpp.plugins.xep_0384.stanza import OMEMO_DEVICES_NS, OMEMO_BUNDLE_NS
+from slixmpp.plugins.xep_0384.stanza import Devices, Device
 from slixmpp.plugins.base import BasePlugin, register_plugin
 
 log = logging.getLogger(__name__)
@@ -60,33 +62,66 @@ class XEP_0384(BasePlugin):
         )
         self._device_id = self._omemo.get_own_device_id()
 
-        self.xmpp.add_event_handler('pubsub_publish', self.device_list)
+        self.xmpp.add_event_handler('pubsub_publish', self._get_device_list)
+
+        asyncio.ensure_future(self._set_device_list())
 
     def plugin_end(self):
         if not self.backend_loaded:
             return
 
-        self.xmpp.del_event_handler('pubsub_publish', self.device_list)
+        self.xmpp.del_event_handler('pubsub_publish', self._get_device_list)
         self.xmpp['xep_0163'].remove_interest(OMEMO_DEVICES_NS)
 
     def session_bind(self, _jid):
         self.xmpp['xep_0163'].add_interest(OMEMO_DEVICES_NS)
 
-    def device_list(self, msg):
+    def _store_device_ids(self, jid, items):
+        self.device_ids[jid] = []
+        for item in items:
+            device_ids = [int(d['id']) for d in item['devices']]
+            self.device_ids[jid] = device_ids
+
+            # XXX: There should only be one item so this is fine, but slixmpp
+            # loops forever otherwise. ???
+            break
+
+    def _get_device_list(self, msg):
         if msg['pubsub_event']['items']['node'] != OMEMO_DEVICES_NS:
             return
 
         jid = msg['from'].bare
         items = msg['pubsub_event']['items']
-        for item in items:
-            device_ids = [d['id'] for d in item['devices']]
-            if jid not in self.device_ids:
-                self.device_ids[jid] = device_ids
-            self.xmpp.event('omemo_device_ids', (jid, device_ids))
+        self._store_device_ids(jid, items)
 
-            # XXX: There should only be one item so this is fine, but slixmpp
-            # loops forever otherwise. ???
+        if jid == self.xmpp.boundjid.bare and \
+           self._device_id not in self.device_ids[jid]:
+            asyncio.ensure_future(self._set_device_list())
+
+    async def _set_device_list(self):
+        iq = await self.xmpp['xep_0060'].get_items(
+            self.xmpp.boundjid.bare, OMEMO_DEVICES_NS,
+        )
+        jid = self.xmpp.boundjid.bare
+        items = iq['pubsub']['items']
+        self._store_device_ids(jid, items)
+
+        # Verify that this device in the list and set it if necessary
+        if self._device_id in self.device_ids[jid]:
             return
 
+        self.device_ids[jid].append(self._device_id)
+
+        devices = []
+        for i in self.device_ids[jid]:
+            d = Device()
+            d['id'] = str(i)
+            devices.append(d)
+        payload = Devices()
+        payload['devices'] = devices
+
+        await self.xmpp['xep_0060'].publish(
+            jid, OMEMO_DEVICES_NS, payload=payload,
+        )
 
 register_plugin(XEP_0384)