From ad5822b3604dc37d3e5cd5bceaf1369a060cdb57 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maxime=20=E2=80=9Cpep=E2=80=9D=20Buquet?= <pep@bouah.net>
Date: Tue, 13 Jul 2021 22:43:40 +0200
Subject: [PATCH] Automatically send heartbeats
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Maxime “pep” Buquet <pep@bouah.net>
---
 slixmpp_omemo/__init__.py | 57 ++++++++++++++++++++++++++++++++++-----
 1 file changed, 50 insertions(+), 7 deletions(-)

diff --git a/slixmpp_omemo/__init__.py b/slixmpp_omemo/__init__.py
index 6c803d0..9d8074b 100644
--- a/slixmpp_omemo/__init__.py
+++ b/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