From 4bdea251798f90639b5c518f86a17d4b7fcb8d38 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, 25 Nov 2018 22:17:47 +0000
Subject: [PATCH] xep_0384: encryption, return empty payload (doesn't encrypt
 anything)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Maxime “pep” Buquet <pep@bouah.net>
---
 plugin.py | 47 +++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 45 insertions(+), 2 deletions(-)

diff --git a/plugin.py b/plugin.py
index ca55dda..c76718b 100644
--- a/plugin.py
+++ b/plugin.py
@@ -8,7 +8,7 @@
 
 import logging
 
-from typing import List, Union
+from typing import Dict, List, Union
 
 import os
 import json
@@ -16,10 +16,11 @@ 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_BUNDLES_NS
-from slixmpp.plugins.xep_0384.stanza import Devices, Device, Key, PreKeyPublic
+from slixmpp.plugins.xep_0384.stanza import Devices, Device, Encrypted, Key, PreKeyPublic
 from slixmpp.plugins.base import BasePlugin, register_plugin
 from slixmpp.exceptions import IqError
 from slixmpp.stanza import Message, Iq
+from slixmpp.jid import JID
 
 log = logging.getLogger(__name__)
 
@@ -253,5 +254,47 @@ class XEP_0384(BasePlugin):
         )
         return body
 
+    async def encrypt_message(self, plaintext: str, recipients: List[JID]) -> Encrypted:
+        """
+        Returns an encrypted payload to be placed into a message.
+
+        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.
+        """
+
+        recipients = [jid.bare for jid in recipients]
+        bundles = {}  # type: Dict[str, Dict[str, ExtendedPublicBundle]]
+
+        while True:
+            errors = []
+
+            self._omemo.encryptMessage(
+                recipients,
+                plaintext.encode('utf-8'),
+                bundles,
+                callback=lambda *args: errors.append(args),
+                always_trust=True,
+                dry_run=True,
+            )
+
+            if not errors:
+                break
+
+            for (exn, key, val) in errors:
+                pass
+
+            break
+
+        # Attempt encryption
+        payload = Encrypted()
+        payload['omemo_encrypted'] = self._omemo.encryptMessage(
+            recipients,
+            plaintext.encode('utf-8'),
+            bundles,
+            always_trust=True,
+        )
+
+        return payload
 
 register_plugin(XEP_0384)