diff sat/plugins/plugin_xep_0384.py @ 3972:5fbdf986670c

plugin pte: Pubsub Target Encryption implementation: This plugin lets encrypt a few items for a specific set of entities. rel 382
author Goffi <goffi@goffi.org>
date Mon, 31 Oct 2022 13:46:51 +0100
parents 8e7d5796fb23
children db45d49518f6
line wrap: on
line diff
--- a/sat/plugins/plugin_xep_0384.py	Mon Oct 31 04:09:38 2022 +0100
+++ b/sat/plugins/plugin_xep_0384.py	Mon Oct 31 13:46:51 2022 +0100
@@ -22,7 +22,7 @@
 import logging
 import time
 from typing import \
-    Any, Dict, FrozenSet, List, Literal, NamedTuple, Optional, Set, Type, cast
+    Any, Dict, FrozenSet, List, Literal, NamedTuple, Optional, Set, Type, Union, cast
 import uuid
 import xml.etree.ElementTree as ET
 from xml.sax.saxutils import quoteattr
@@ -1581,7 +1581,7 @@
         # added, the messageReceived trigger is also used for twomemo.
         sat.trigger.add(
             "messageReceived",
-            self.__message_received_trigger,
+            self._message_received_trigger,
             priority=100050
         )
         sat.trigger.add(
@@ -2098,7 +2098,7 @@
                 frozenset(applied_trust_updates)
             )
 
-    async def __message_received_trigger(
+    async def _message_received_trigger(
         self,
         client: SatXMPPClient,
         message_elt: domish.Element,
@@ -2113,13 +2113,12 @@
             encrypted.
         @return: Whether to continue the message received flow.
         """
-
         muc_plaintext_cache_key: Optional[MUCPlaintextCacheKey] = None
 
         sender_jid = jid.JID(message_elt["from"])
         feedback_jid: jid.JID
 
-        message_type = message_elt.getAttribute("type", "unknown")
+        message_type = message_elt.getAttribute("type", C.MESS_TYPE_NORMAL)
         is_muc_message = message_type == C.MESS_TYPE_GROUPCHAT
         if is_muc_message:
             if self.__xep_0045 is None:
@@ -2420,12 +2419,12 @@
             return True
 
         # All pre-checks done, we can start encrypting!
-        await self.__encrypt(
+        await self.encrypt(
             client,
             twomemo.twomemo.NAMESPACE,
             stanza,
             recipient_bare_jid,
-            stanza.getAttribute("type", "unkown") == C.MESS_TYPE_GROUPCHAT,
+            stanza.getAttribute("type", C.MESS_TYPE_NORMAL) == C.MESS_TYPE_GROUPCHAT,
             stanza.getAttribute("id", None)
         )
 
@@ -2462,7 +2461,7 @@
         is_muc_message = mess_data["type"] == C.MESS_TYPE_GROUPCHAT
         stanza_id = mess_data["uid"]
 
-        await self.__encrypt(
+        await self.encrypt(
             client,
             oldmemo.oldmemo.NAMESPACE,
             stanza,
@@ -2474,12 +2473,12 @@
         # Add a store hint
         self.__xep_0334.addHintElements(stanza, [ "store" ])
 
-    async def __encrypt(
+    async def encrypt(
         self,
         client: SatXMPPClient,
         namespace: Literal["urn:xmpp:omemo:2", "eu.siacs.conversations.axolotl"],
         stanza: domish.Element,
-        recipient_jid: jid.JID,
+        recipient_jids: Union[jid.JID, Set[jid.JID]],
         is_muc_message: bool,
         stanza_id: Optional[str]
     ) -> None:
@@ -2488,8 +2487,9 @@
         @param namespace: The namespace of the OMEMO version to use.
         @param stanza: The stanza. Twomemo will encrypt the whole stanza using SCE,
             oldmemo will encrypt only the body. The stanza is modified by this call.
-        @param recipient_jid: The JID of the recipient. Can be a bare (aka "userhost") JID
-            but doesn't have to.
+        @param recipient_jid: The JID of the recipients.
+            Can be a bare (aka "userhost") JIDs but doesn't have to.
+            A single JID can be used.
         @param is_muc_message: Whether the stanza is a message stanza to a MUC room.
         @param stanza_id: The id of this stanza. Especially relevant for message stanzas
             to MUC rooms such that the outgoing plaintext can be cached for MUC message
@@ -2499,6 +2499,11 @@
             hint to the stanza if applicable! This can be done before or after this call,
             the order doesn't matter.
         """
+        if isinstance(recipient_jids, jid.JID):
+            recipient_jids = {recipient_jids}
+        if not recipient_jids:
+            raise exceptions.InternalError("At least one JID must be specified")
+        recipient_jid = next(iter(recipient_jids))
 
         muc_plaintext_cache_key: Optional[MUCPlaintextCacheKey] = None
 
@@ -2506,6 +2511,10 @@
         feedback_jid: jid.JID
 
         if is_muc_message:
+            if len(recipient_jids) != 1:
+                raise exceptions.InternalError(
+                    'Only one JID can be set when "is_muc_message" is set'
+                )
             if self.__xep_0045 is None:
                 raise exceptions.InternalError(
                     "Encryption of MUC message requested, but plugin XEP-0045 is not"
@@ -2531,7 +2540,7 @@
                 message_uid=stanza_id
             )
         else:
-            recipient_bare_jids = { recipient_jid.userhost() }
+            recipient_bare_jids = {r.userhost() for r in recipient_jids}
             feedback_jid = recipient_jid.userhostJID()
 
         log.debug(