From 0c38ee374c3bfc92c81fc3dde7183b145736bfa2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philipp=20H=C3=B6rist?= <philipp@hoerist.com>
Date: Sun, 24 Feb 2019 12:26:45 +0100
Subject: [PATCH] [omemo] Make sending key transport messages possible

---
 omemo/backend/aes.py   |  8 ++++++++
 omemo/backend/state.py | 25 ++++++++++++++++++++++++-
 omemo/modules/omemo.py | 11 +++++++++++
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/omemo/backend/aes.py b/omemo/backend/aes.py
index 581697a..eec2b7c 100644
--- a/omemo/backend/aes.py
+++ b/omemo/backend/aes.py
@@ -84,3 +84,11 @@ def aes_encrypt_file(data):
     key, iv, tag, payload, = _encrypt(data, 32, 16)
     payload += tag
     return EncryptionResult(payload=payload, key=key, iv=iv)
+
+
+def get_new_key():
+    return os.urandom(16)
+
+
+def get_new_iv():
+    return os.urandom(16)
diff --git a/omemo/backend/state.py b/omemo/backend/state.py
index 5e4237a..a881785 100644
--- a/omemo/backend/state.py
+++ b/omemo/backend/state.py
@@ -33,7 +33,10 @@ from axolotl.state.prekeybundle import PreKeyBundle
 from axolotl.util.keyhelper import KeyHelper
 from axolotl.duplicatemessagexception import DuplicateMessageException
 
-from omemo.backend.aes import aes_decrypt, aes_encrypt
+from omemo.backend.aes import aes_decrypt
+from omemo.backend.aes import aes_encrypt
+from omemo.backend.aes import get_new_key
+from omemo.backend.aes import get_new_iv
 from omemo.backend.devices import DeviceManager
 from omemo.backend.devices import NoDevicesFound
 from omemo.backend.liteaxolotlstore import LiteAxolotlStore
@@ -201,6 +204,26 @@ class OmemoState(DeviceManager):
                             iv=result.iv,
                             payload=result.payload)
 
+    def encrypt_key_transport(self, jid, devices):
+        whisper_messages = defaultdict(dict)
+        for device in devices:
+            try:
+                whisper_messages[jid][device] = self._get_whisper_message(
+                    jid, device, get_new_key())
+            except Exception:
+                log.exception('Failed to encrypt')
+                continue
+
+        if not whisper_messages[jid]:
+            log.error('Encrypted keys empty')
+            return
+
+        log.debug('Finished Key Transport message')
+        return OMEMOMessage(sid=self.own_device,
+                            keys=whisper_messages[jid],
+                            iv=get_new_iv(),
+                            payload=None)
+
     def has_trusted_keys(self, jid):
         inactive = self._storage.getInactiveSessionsKeys(jid)
         trusted = self._storage.getTrustedFingerprints(jid)
diff --git a/omemo/modules/omemo.py b/omemo/modules/omemo.py
index 91c8e3a..1aadd38 100644
--- a/omemo/modules/omemo.py
+++ b/omemo/modules/omemo.py
@@ -29,6 +29,7 @@ from nbxmpp.const import PresenceType
 from nbxmpp.const import Affiliation
 from nbxmpp.structs import StanzaHandler
 from nbxmpp.modules.omemo import create_omemo_message
+from nbxmpp.modules.omemo import get_key_transport_message
 
 from gajim.common import app
 from gajim.common import helpers
@@ -192,6 +193,16 @@ class OMEMO(BaseModule):
         self._debug_print_stanza(event.msg_iq)
         callback(event)
 
+    def _send_key_transport_message(self, typ, jid, devices):
+        omemo_message = self.backend.encrypt_key_transport(jid, devices)
+        if omemo_message is None:
+            log.warning('Key transport message to %s (%s) failed', jid, devices)
+            return
+
+        transport_message = get_key_transport_message(typ, jid, omemo_message)
+        log.info('Send key transport message %s (%s)', jid, devices)
+        self._con.connection.send(transport_message)
+
     def _message_received(self, _con, stanza, properties):
         if not properties.is_omemo:
             return