# HG changeset patch
# User Goffi <goffi@goffi.org>
# Date 1709656296 -3600
# Node ID 1b5cf2ee1d86e5bd3f3edfb15e2f7bc779c48fe2
# Parent  c0f3f29377f1291ff332affb65ecb42fa06c38da
plugin XEP-0384, XEP-0391: download missing devices list:

when a peer jid was not in our roster, devices list was not retrieved, resulting in failed
en/decryption. This patch does check it and download missing devices list in necessary.

There is no subscription managed yet, so the list won't be updated in case of new devices,
this should be addressed at some point.

diff -r c0f3f29377f1 -r 1b5cf2ee1d86 libervia/backend/plugins/plugin_xep_0384.py
--- a/libervia/backend/plugins/plugin_xep_0384.py	Tue Mar 05 17:31:36 2024 +0100
+++ b/libervia/backend/plugins/plugin_xep_0384.py	Tue Mar 05 17:31:36 2024 +0100
@@ -22,7 +22,7 @@
 import logging
 import time
 from typing import \
-    Any, Dict, FrozenSet, List, Literal, NamedTuple, Optional, Set, Type, Union, cast
+    Any, Dict, FrozenSet, Iterable, List, Literal, NamedTuple, Optional, Set, Type, Union, cast
 import uuid
 import xml.etree.ElementTree as ET
 from xml.sax.saxutils import quoteattr
@@ -97,6 +97,8 @@
 PARAM_CATEGORY = "Security"
 PARAM_NAME = "omemo_policy"
 
+NamespaceType = Literal["urn:xmpp:omemo:2", "eu.siacs.conversations.axolotl"]
+
 
 class LogHandler(logging.Handler):
     """
@@ -2479,10 +2481,50 @@
         # Let the flow continue.
         return True
 
+    async def download_missing_device_lists(
+        self,
+        client: SatXMPPClient,
+        namespace: NamespaceType,
+        recipients: Iterable[jid.JID],
+        session_manager: omemo.SessionManager,
+    ) -> None:
+        """Retrieves missing device lists for recipients outside the profile's roster.
+
+        @param client: XMPP client.
+        @param namespace: The namespace of the OMEMO version to use.
+        @param recipients: Recipients to verify device list presence.
+        @param session_manager: OMEMO session manager.
+        """
+        recipients = [j.userhostJID() for j in recipients]
+        not_in_roster = [j for j in recipients if not client.roster.is_jid_in_roster(j)]
+        for bare_jid in not_in_roster:
+            device_information = await session_manager.get_device_information(
+                bare_jid.userhost()
+            )
+            if (
+                not device_information
+                or not all(namespace in di.namespaces for di in device_information)
+            ):
+                if namespace == self.NS_TWOMEMO:
+                    algo, node = "OMEMO", TWOMEMO_DEVICE_LIST_NODE
+                elif namespace == self.NS_OLDMEMO:
+                    algo, node = "OMEMO_legacy", OLDMEMO_DEVICE_LIST_NODE
+                else:
+                    raise ValueError(f"Invalid namespace: {namespace!r}")
+
+                try:
+                    items, __ = await self._j.get_items(client, bare_jid, node, 1)
+
+                except Exception:
+                    log.exception(f"Can't find {algo} devices list for {bare_jid}.")
+                else:
+                    await self._update_device_list(client, bare_jid, items)
+                    log.warning(f"{algo} devices list updated for {bare_jid}.")
+
     async def encrypt(
         self,
         client: SatXMPPClient,
-        namespace: Literal["urn:xmpp:omemo:2", "eu.siacs.conversations.axolotl"],
+        namespace: NamespaceType,
         stanza: domish.Element,
         recipient_jids: Union[jid.JID, Set[jid.JID]],
         is_muc_message: bool,
@@ -2602,6 +2644,7 @@
         log.debug(f"Plaintext to encrypt: {plaintext}")
 
         session_manager = await self.get_session_manager(client.profile)
+        await self.download_missing_device_lists(client, namespace, recipient_jids, session_manager)
 
         try:
             messages, encryption_errors = await session_manager.encrypt(
@@ -2679,6 +2722,15 @@
 
         sender = cast(jid.JID, items_event.sender)
         items = cast(List[domish.Element], items_event.items)
+        client = self.host.get_client(profile)
+        await self._update_device_list(client, sender, items)
+
+    async def _update_device_list(
+        self,
+        client: SatXMPPEntity,
+        sender: jid.JID,
+        items: list[domish.Element]
+    ) -> None:
 
         if len(items) > 1:
             log.warning("Ignoring device list update with more than one element.")
@@ -2719,7 +2771,7 @@
             )
             return
 
-        session_manager = await self.get_session_manager(profile)
+        session_manager = await self.get_session_manager(client.profile)
 
         await session_manager.update_device_list(
             namespace,
diff -r c0f3f29377f1 -r 1b5cf2ee1d86 libervia/backend/plugins/plugin_xep_0391.py
--- a/libervia/backend/plugins/plugin_xep_0391.py	Tue Mar 05 17:31:36 2024 +0100
+++ b/libervia/backend/plugins/plugin_xep_0391.py	Tue Mar 05 17:31:36 2024 +0100
@@ -124,6 +124,9 @@
                 "type": enc_type
             }
             session_manager = await self._o.get_session_manager(client.profile)
+            await self._o.download_missing_device_lists(
+                client, self._o.NS_OLDMEMO, {session["peer_jid"]}, session_manager
+            )
             try:
                 messages, encryption_errors = await session_manager.encrypt(
                     frozenset({session["peer_jid"].userhost()}),
@@ -133,7 +136,7 @@
                     identifier = client.jid.userhost()
                 )
             except Exception as e:
-                log.error("Can't generate IV and keys: {e}")
+                log.exception("Can't generate IV and keys")
                 raise e
             message, plain_key_material = next(iter(messages.items()))
             iv, key = message.content.initialization_vector, plain_key_material.key