Automatically send heartbeats

Signed-off-by: Maxime “pep” Buquet <pep@bouah.net>

slixmpp_omemo/__init__.py

@@ -117,6 +117,7 @@ def _generate_encrypted_payload(encrypted) -> Encrypted:
 
     tag['header']['sid'] = str(encrypted['sid'])
     tag['header']['iv']['value'] = b64enc(encrypted['iv'])
+    if 'payload' in encrypted:
         tag['payload']['value'] = b64enc(encrypted['payload'])
 
     for bare_jid, devices in encrypted['keys'].items():
@@ -197,6 +198,7 @@ class XEP_0384(BasePlugin):
         'storage_backend': None,
         'otpk_policy': DefaultOTPKPolicy,
         'omemo_backend': SignalBackend,
+        'auto_heartbeat': True,
         'heartbeat_after': 53,
         # TODO: 'drop_inactive_after': 300,
     }
@@ -536,6 +538,24 @@ class XEP_0384(BasePlugin):
         lengths = map(lambda d_l: d_l[1], receiving_chain_lengths)
         return max(lengths, default=0) > self.heartbeat_after
 
+    def make_heartbeat(self, jid: JID) -> Message:
+        """
+            Returns a heartbeat message.
+
+            This is mainly used to tell receiving clients that our device is
+            still active. This is an empty key transport message of which we
+            won't use the generated shared secret.
+        """
+
+        msg = self.xmpp.make_message(jid)
+        encrypted = self.encrypt_key_transport_message(
+            plaintext=None,
+            recipients=[jid],
+            expect_problems=None,
+        )
+        msg.append(encrypted)
+        return msg
+
     def trust(self, jid: JID, device_id: int, ik: bytes) -> None:
         self._omemo.setTrust(jid.bare, device_id, ik, True)
 
@@ -623,6 +643,10 @@ class XEP_0384(BasePlugin):
         finally:
             asyncio.ensure_future(self._publish_bundle())
 
+        if self.auto_heartbeat:
+            msg = self.make_heartbeat(jid)
+            asyncio.ensure_future(msg.send())
+
         return body
 
     async def encrypt_message(
@@ -630,6 +654,18 @@ class XEP_0384(BasePlugin):
         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],
+        recipients: List[JID],
+        expect_problems: Optional[Dict[JID, List[int]]] = None,
     ) -> Encrypted:
         """
         Returns an encrypted payload to be placed into a message.
@@ -653,9 +689,16 @@ class XEP_0384(BasePlugin):
             expect_problems = {jid.bare: did for (jid, did) in expect_problems.items()}
 
             try:
+                if plaintext is not None:
                     encrypted = self._omemo.encryptMessage(
                         recipients,
-                    plaintext.encode('utf-8'),
+                        plaintext,
+                        self.bundles,
+                        expect_problems=expect_problems,
+                    )
+                else:
+                    encrypted = self._omemo.encryptKeyTransportMessage(
+                        recipients,
                         self.bundles,
                         expect_problems=expect_problems,
                     )