From f20e5d610392678396bdc28b247c233c80cae8ab Mon Sep 17 00:00:00 2001
From: Goffi <goffi@goffi.org>
Date: Tue, 5 Mar 2024 17:31:36 +0100
Subject: [PATCH] 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.
---
 libervia/backend/plugins/plugin_xep_0384.py | 58 +++++++++++++++++++--
 libervia/backend/plugins/plugin_xep_0391.py |  5 +-
 2 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/libervia/backend/plugins/plugin_xep_0384.py b/libervia/backend/plugins/plugin_xep_0384.py
index a6990987..3c1dbdf1 100644
--- a/libervia/backend/plugins/plugin_xep_0384.py
+++ b/libervia/backend/plugins/plugin_xep_0384.py
@@ -22,7 +22,7 @@ import enum
 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 @@ PLUGIN_INFO = {
 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 @@ class OMEMO:
         # 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 @@ class OMEMO:
         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 @@ class OMEMO:
 
         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 @@ class OMEMO:
             )
             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 --git a/libervia/backend/plugins/plugin_xep_0391.py b/libervia/backend/plugins/plugin_xep_0391.py
index 45513a03..5acb7b75 100644
--- a/libervia/backend/plugins/plugin_xep_0391.py
+++ b/libervia/backend/plugins/plugin_xep_0391.py
@@ -124,6 +124,9 @@ class JET:
                 "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 @@ class JET:
                     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