[omemo] Set prekey attribute in PreKeyMessage

2 files changed, 20 insertions, 8 deletions

diff --git a/omemo/omemo/state.py b/omemo/omemo/state.py
index c0aeb93..58c40ca 100644
--- a/omemo/omemo/state.py
+++ b/omemo/omemo/state.py
@@ -237,7 +237,9 @@ class OmemoState:
             try:
                 if self.isTrusted(jid, device) == TRUSTED:
                     cipher = self.get_session_cipher(jid, device)
-                    encrypted_keys[device] = cipher.encrypt(key).serialize()
+                    cipher_key = cipher.encrypt(key)
+                    prekey = isinstance(cipher_key, PreKeyWhisperMessage)
+                    encrypted_keys[device] = (cipher_key.serialize(), prekey)
                 else:
                     log.debug('Skipped Device because Trust is: ' +
                               str(self.isTrusted(jid, device)))
@@ -254,7 +256,9 @@ class OmemoState:
             try:
                 if self.isTrusted(from_jid, device) == TRUSTED:
                     cipher = self.get_session_cipher(from_jid, device)
-                    encrypted_keys[device] = cipher.encrypt(key).serialize()
+                    cipher_key = cipher.encrypt(key)
+                    prekey = isinstance(cipher_key, PreKeyWhisperMessage)
+                    encrypted_keys[device] = (cipher_key.serialize(), prekey)
                 else:
                     log.debug('Skipped own Device because Trust is: ' +
                               str(self.isTrusted(from_jid, device)))
@@ -302,8 +306,9 @@ class OmemoState:
             for rid, cipher in self.session_ciphers[jid_to].items():
                 try:
                     if self.isTrusted(jid_to, rid) == TRUSTED:
-                        encrypted_keys[rid] = cipher.encrypt(key). \
-                            serialize()
+                        cipher_key = cipher.encrypt(key)
+                        prekey = isinstance(cipher_key, PreKeyWhisperMessage)
+                        encrypted_keys[rid] = (cipher_key.serialize(), prekey)
                     else:
                         log.debug('Skipped Device because Trust is: ' +
                                   str(self.isTrusted(jid_to, rid)))
@@ -323,7 +328,9 @@ class OmemoState:
             try:
                 cipher = self.get_session_cipher(from_jid, dev)
                 if self.isTrusted(from_jid, dev) == TRUSTED:
-                    encrypted_keys[dev] = cipher.encrypt(key).serialize()
+                    cipher_key = cipher.encrypt(key)
+                    prekey = isinstance(cipher_key, PreKeyWhisperMessage)
+                    encrypted_keys[dev] = (cipher_key.serialize(), prekey)
                 else:
                     log.debug('Skipped own Device because Trust is: ' +
                               str(self.isTrusted(from_jid, dev)))
diff --git a/omemo/xmpp.py b/omemo/xmpp.py
index bf31c8d..b7a6696 100644
--- a/omemo/xmpp.py
+++ b/omemo/xmpp.py
@@ -79,9 +79,14 @@ class OmemoMessage(Node):
         # , contact_jid, key, iv, payload, dev_id, my_dev_id):
         Node.__init__(self, 'encrypted', attrs={'xmlns': NS_OMEMO})
         header = Node('header', attrs={'sid': msg_dict['sid']})
-        for rid, key in msg_dict['keys'].items():
-            header.addChild('key', attrs={'rid': rid}).addData(b64encode(key))
-
+        for rid, (key, prekey) in msg_dict['keys'].items():
+            if prekey:
+                child = header.addChild('key',
+                                        attrs={'prekey': 'true', 'rid': rid})
+            else:
+                child = header.addChild('key',
+                                        attrs={'rid': rid})
+            child.addData(b64encode(key))
         header.addChild('iv').addData(b64encode(msg_dict['iv']))
         self.addChild(node=header)
         self.addChild('payload').addData(b64encode(msg_dict['payload']))