diff --git a/slixmpp_omemo/__init__.py b/slixmpp_omemo/__init__.py index 6c803d0..b31d238 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 + async 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(mto=jid) + encrypted = await 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,12 @@ class XEP_0384(BasePlugin): finally: asyncio.ensure_future(self._publish_bundle()) + if self.auto_heartbeat and self.should_heartbeat(): + async def send_heartbeat(): + msg = await self.make_heartbeat(JID(jid)) + msg.send() + asyncio.ensure_future(send_heartbeat()) + return body async def encrypt_message( @@ -630,6 +656,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 +691,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