From 3a7711aaf1e6e81d787a534163c5446de8433a2a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maxime=20=E2=80=9Cpep=E2=80=9D=20Buquet?= <pep@bouah.net>
Date: Thu, 31 Jan 2019 15:47:12 +0100
Subject: [PATCH] xep_0384: when encrypting, only tell client when trust is
 undecided
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

On UntrustedException, there are two possibilities.  Either trust has not been
explicitely set yet, and is 'undecided', or the device is explicitely not
trusted. When undecided, we need to ask our user to make a choice. If
untrusted, then we can safely tell the OMEMO lib to not encrypt to this device.

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

diff --git a/plugin.py b/plugin.py
index 5bf19a0..d6858d4 100644
--- a/plugin.py
+++ b/plugin.py
@@ -119,6 +119,9 @@ class EncryptionPrepareException(XEP0384): pass
 class UntrustedException(XEP0384): pass
 
 
+class UndecidedException(XEP0384): pass
+
+
 class XEP_0384(BasePlugin):
 
     """
@@ -346,7 +349,12 @@ class XEP_0384(BasePlugin):
         finally:
             asyncio.ensure_future(self._publish_bundle())
 
-    async def encrypt_message(self, plaintext: str, recipients: List[JID]) -> Encrypted:
+    async def encrypt_message(
+        self,
+        plaintext: str,
+        recipients: List[JID],
+        expect_problems: Optional[Dict[JID, List[int]]] = None,
+    ) -> Encrypted:
         """
         Returns an encrypted payload to be placed into a message.
 
@@ -364,11 +372,15 @@ class XEP_0384(BasePlugin):
             # or if we hit the same set of errors.
             errors = []  # type: List[omemo.exceptions.OMEMOException]
 
+            if expect_problems is not None:
+                expect_problems = {jid.bare: did for (jid, did) in expect_problems.items()}
+
             try:
                 encrypted = self._omemo.encryptMessage(
                     recipients,
                     plaintext.encode('utf-8'),
                     bundles,
+                    expect_problems=expect_problems,
                 )
                 return _generate_encrypted_payload(encrypted)
             except omemo.exceptions.EncryptionProblemsException as e:
@@ -389,7 +401,15 @@ class XEP_0384(BasePlugin):
                         devices = bundles.setdefault(exn.bare_jid, {})
                         devices[exn.device] = bundle
                 elif isinstance(exn, omemo.exceptions.UntrustedException):
-                    raise UntrustedException(exn.bare_jid, exn.device, exn.ik)
+                    # On UntrustedException, there are two possibilities.
+                    # Either trust has not been explicitely set yet, and is
+                    # 'undecided', or the device is explicitely not
+                    # trusted. When undecided, we need to ask our user to make
+                    # a choice. If untrusted, then we can safely tell the
+                    # OMEMO lib to not encrypt to this device
+                    if self._omemo.getTrustForDevice(exn.bare_jid, exn.device) is None:
+                        raise UndecidedException(exn.bare_jid, exn.device, exn.ik)
+                    expect_problems.setdefault(exn.bare_jid, []).append(exn.device)
                 elif isinstance(exn, omemo.exceptions.NoEligibleDevicesException):
                     # This error is returned by the library to specify that
                     # encryption is not possible to any device of a user.