From 494899bb3c0d4a6ef90d96babb82b82220fd0825 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maxime=20=E2=80=9Cpep=E2=80=9D=20Buquet?= <pep@bouah.net>
Date: Sun, 18 Jul 2021 00:30:07 +0200
Subject: [PATCH] Add interface semi-hidden to RatchetForwardingMessage
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Maxime “pep” Buquet <pep@bouah.net>
---
 slixmpp_omemo/__init__.py | 34 +++++++++++++++++++---------------
 1 file changed, 19 insertions(+), 15 deletions(-)

diff --git a/slixmpp_omemo/__init__.py b/slixmpp_omemo/__init__.py
index dafd807..e5381bd 100644
--- a/slixmpp_omemo/__init__.py
+++ b/slixmpp_omemo/__init__.py
@@ -602,10 +602,11 @@ class XEP_0384(BasePlugin):
         """
 
         msg = self.xmpp.make_message(mto=jid)
-        encrypted = await self.encrypt_key_transport_message(
+        encrypted = await self.encrypt_message(
             plaintext=None,
             recipients=[jid],
             expect_problems=None,
+            _ignore_trust=True,
         )
         msg.append(encrypted)
         return msg
@@ -708,21 +709,10 @@ class XEP_0384(BasePlugin):
 
     async def encrypt_message(
         self,
-        plaintext: str,
-        recipients: List[JID],
-        expect_problems: Optional[Dict[JID, List[int]]] = None,
-    ) -> Encrypted:
-        return await self.encrypt_key_transport_message(
-            plaintext.encode('utf-8'),
-            recipients,
-            expect_problems,
-        )
-
-    async def encrypt_key_transport_message(
-        self,
-        plaintext: Optional[bytes],
+        plaintext: Optional[str],
         recipients: List[JID],
         expect_problems: Optional[Dict[JID, List[int]]] = None,
+        _ignore_trust: bool = False,
     ) -> Encrypted:
         """
         Returns an encrypted payload to be placed into a message.
@@ -730,6 +720,14 @@ class XEP_0384(BasePlugin):
         The API for getting an encrypted payload consists of trying first
         and fixing errors progressively. The actual sending happens once the
         application (us) thinks we're good to go.
+
+        If `plaintext` is specified, this will generate a full OMEMO payload. If
+        not, if `_ignore_trust` is True, this will generate a ratchet forwarding
+        message, and otherwise it will generate a key transport message.
+
+        These are rather technical details to the user and fiddling with
+        parameters else than `plaintext` and `recipients` should be rarely
+        needed.
         """
 
         recipients = [jid.bare for jid in recipients]
@@ -749,7 +747,13 @@ class XEP_0384(BasePlugin):
                 if plaintext is not None:
                     encrypted = self._omemo.encryptMessage(
                         recipients,
-                        plaintext,
+                        plaintext.encode('utf-8'),
+                        self.bundles,
+                        expect_problems=expect_problems,
+                    )
+                elif _ignore_trust:
+                    encrypted = self._omemo.encryptRatchetForwardingMessage(
+                        recipients,
                         self.bundles,
                         expect_problems=expect_problems,
                     )