Automatically send heartbeats

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

slixmpp_omemo/__init__.py

@@ -117,7 +117,8 @@ def _generate_encrypted_payload(encrypted) -> Encrypted:
 
     tag['header']['sid'] = str(encrypted['sid'])
     tag['header']['iv']['value'] = b64enc(encrypted['iv'])
-    tag['payload']['value'] = b64enc(encrypted['payload'])
+    if 'payload' in encrypted:
+        tag['payload']['value'] = b64enc(encrypted['payload'])
 
     for bare_jid, devices in encrypted['keys'].items():
         for rid, device in devices.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,12 +689,19 @@ class XEP_0384(BasePlugin):
             expect_problems = {jid.bare: did for (jid, did) in expect_problems.items()}
 
             try:
-                encrypted = self._omemo.encryptMessage(
-                    recipients,
-                    plaintext.encode('utf-8'),
-                    self.bundles,
-                    expect_problems=expect_problems,
-                )
+                if plaintext is not None:
+                    encrypted = self._omemo.encryptMessage(
+                        recipients,
+                        plaintext,
+                        self.bundles,
+                        expect_problems=expect_problems,
+                    )
+                else:
+                    encrypted = self._omemo.encryptKeyTransportMessage(
+                        recipients,
+                        self.bundles,
+                        expect_problems=expect_problems,
+                    )
                 return _generate_encrypted_payload(encrypted)
             except omemo.exceptions.EncryptionProblemsException as exception:
                 errors = exception.problems