From ad5822b3604dc37d3e5cd5bceaf1369a060cdb57 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maxime=20=E2=80=9Cpep=E2=80=9D=20Buquet?= Date: Tue, 13 Jul 2021 22:43:40 +0200 Subject: [PATCH 1/5] Automatically send heartbeats MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Maxime “pep” Buquet --- 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 From 7f1d48c529bb805de78342569a7d20b9bda02f95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maxime=20=E2=80=9Cpep=E2=80=9D=20Buquet?= Date: Tue, 13 Jul 2021 23:14:04 +0200 Subject: [PATCH 2/5] make_heartbeat: add parameter to make_message MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Maxime “pep” Buquet --- slixmpp_omemo/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/slixmpp_omemo/__init__.py b/slixmpp_omemo/__init__.py index 9d8074b..e60d752 100644 --- a/slixmpp_omemo/__init__.py +++ b/slixmpp_omemo/__init__.py @@ -547,7 +547,7 @@ class XEP_0384(BasePlugin): won't use the generated shared secret. """ - msg = self.xmpp.make_message(jid) + msg = self.xmpp.make_message(mto=jid) encrypted = self.encrypt_key_transport_message( plaintext=None, recipients=[jid], From 89eb4dfecee9a275255975ae713f2c34d51dbd6b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maxime=20=E2=80=9Cpep=E2=80=9D=20Buquet?= Date: Tue, 13 Jul 2021 23:14:33 +0200 Subject: [PATCH 3/5] decrypt_message: msg.send isn't a coroutine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Maxime “pep” Buquet --- slixmpp_omemo/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/slixmpp_omemo/__init__.py b/slixmpp_omemo/__init__.py index e60d752..0ecb0d2 100644 --- a/slixmpp_omemo/__init__.py +++ b/slixmpp_omemo/__init__.py @@ -645,7 +645,7 @@ class XEP_0384(BasePlugin): if self.auto_heartbeat: msg = self.make_heartbeat(jid) - asyncio.ensure_future(msg.send()) + msg.send() return body From 95481e64b2ee226ba97fab70c98a6e284b1af347 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maxime=20=E2=80=9Cpep=E2=80=9D=20Buquet?= Date: Tue, 13 Jul 2021 23:15:02 +0200 Subject: [PATCH 4/5] decrypt_message: don't always send heartbeat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Maxime “pep” Buquet --- slixmpp_omemo/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/slixmpp_omemo/__init__.py b/slixmpp_omemo/__init__.py index 0ecb0d2..3c2f689 100644 --- a/slixmpp_omemo/__init__.py +++ b/slixmpp_omemo/__init__.py @@ -643,7 +643,7 @@ class XEP_0384(BasePlugin): finally: asyncio.ensure_future(self._publish_bundle()) - if self.auto_heartbeat: + if self.auto_heartbeat and self.should_heartbeat(): msg = self.make_heartbeat(jid) msg.send() From 7e079f4260f16c0c07b413e8074a10d59ebea967 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maxime=20=E2=80=9Cpep=E2=80=9D=20Buquet?= Date: Tue, 13 Jul 2021 23:48:39 +0200 Subject: [PATCH 5/5] make_heartbeat needs to be async MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Maxime “pep” Buquet --- slixmpp_omemo/__init__.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/slixmpp_omemo/__init__.py b/slixmpp_omemo/__init__.py index 3c2f689..b31d238 100644 --- a/slixmpp_omemo/__init__.py +++ b/slixmpp_omemo/__init__.py @@ -538,7 +538,7 @@ 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: + async def make_heartbeat(self, jid: JID) -> Message: """ Returns a heartbeat message. @@ -548,7 +548,7 @@ class XEP_0384(BasePlugin): """ msg = self.xmpp.make_message(mto=jid) - encrypted = self.encrypt_key_transport_message( + encrypted = await self.encrypt_key_transport_message( plaintext=None, recipients=[jid], expect_problems=None, @@ -644,8 +644,10 @@ class XEP_0384(BasePlugin): asyncio.ensure_future(self._publish_bundle()) if self.auto_heartbeat and self.should_heartbeat(): - msg = self.make_heartbeat(jid) - msg.send() + async def send_heartbeat(): + msg = await self.make_heartbeat(JID(jid)) + msg.send() + asyncio.ensure_future(send_heartbeat()) return body