From 3a8fb991c8dd226bce4b6e8296c1d08865cc2f01 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philipp=20H=C3=B6rist?= Date: Sat, 16 Feb 2019 10:58:13 +0100 Subject: [PATCH] [omemo] More refactoring --- omemo/backend/state.py | 203 ++++++++++++++++++++++------------------- omemo/modules/omemo.py | 119 +++++++++++++----------- 2 files changed, 176 insertions(+), 146 deletions(-) diff --git a/omemo/backend/state.py b/omemo/backend/state.py index 8a913d3..429835c 100644 --- a/omemo/backend/state.py +++ b/omemo/backend/state.py @@ -17,6 +17,9 @@ import logging import time +import binascii +import textwrap +from collections import defaultdict from nbxmpp.structs import OMEMOBundle from nbxmpp.structs import OMEMOMessage @@ -31,6 +34,7 @@ from axolotl.sessionbuilder import SessionBuilder from axolotl.sessioncipher import SessionCipher from axolotl.state.prekeybundle import PreKeyBundle from axolotl.util.keyhelper import KeyHelper +from axolotl.duplicatemessagexception import DuplicateMessageException from omemo.backend.aes import aes_decrypt, aes_encrypt from omemo.backend.liteaxolotlstore import LiteAxolotlStore @@ -52,7 +56,7 @@ class OmemoState: def __init__(self, own_jid, db_con, account, xmpp_con): self.account = account self.xmpp_con = xmpp_con - self.session_ciphers = {} + self._session_ciphers = defaultdict(dict) self.own_jid = own_jid self.device_ids = {} self.own_devices = [] @@ -71,40 +75,31 @@ class OmemoState: self.account, self.store.getPreKeyCount()) - def build_session(self, recipient_id, device_id, bundle): - sessionBuilder = SessionBuilder(self.store, self.store, self.store, - self.store, recipient_id, device_id) + def build_session(self, jid, device_id, bundle): + session = SessionBuilder(self.store, self.store, self.store, + self.store, jid, device_id) registration_id = self.store.getLocalRegistrationId() prekey = bundle.pick_prekey() - preKeyPublic = DjbECPublicKey(prekey['key'][1:]) + otpk = DjbECPublicKey(prekey['key'][1:]) - signedPreKeyPublic = DjbECPublicKey(bundle.spk['key'][1:]) - identityKey = IdentityKey(DjbECPublicKey(bundle.ik[1:])) + spk = DjbECPublicKey(bundle.spk['key'][1:]) + ik = IdentityKey(DjbECPublicKey(bundle.ik[1:])) - prekey_bundle = PreKeyBundle( - registration_id, device_id, - prekey['id'], preKeyPublic, - bundle.spk['id'], signedPreKeyPublic, - bundle.spk_signature, - identityKey) + prekey_bundle = PreKeyBundle(registration_id, + device_id, + prekey['id'], + otpk, + bundle.spk['id'], + spk, + bundle.spk_signature, + ik) - sessionBuilder.processPreKeyBundle(prekey_bundle) - return self.get_session_cipher(recipient_id, device_id) + session.processPreKeyBundle(prekey_bundle) + return self._get_session_cipher(jid, device_id) def set_devices(self, name, devices): - """ Return a an. - - Parameters - ---------- - jid : string - The contacts jid - - devices: [int] - A list of devices - """ - self.device_ids[name] = devices log.info('%s => Saved devices for %s', self.account, name) @@ -146,62 +141,62 @@ class OmemoState: @property def bundle(self): - self.checkPreKeyAmount() + self._check_pre_key_count() bundle = {'otpks': []} for k in self.store.loadPendingPreKeys(): key = k.getKeyPair().getPublicKey().serialize() bundle['otpks'].append({'key': key, 'id': k.getId()}) - identityKeyPair = self.store.getIdentityKeyPair() - bundle['ik'] = identityKeyPair.getPublicKey().serialize() + ik_pair = self.store.getIdentityKeyPair() + bundle['ik'] = ik_pair.getPublicKey().serialize() - self.cycleSignedPreKey(identityKeyPair) + self._cycle_signed_pre_key(ik_pair) - signedPreKey = self.store.loadSignedPreKey( + spk = self.store.loadSignedPreKey( self.store.getCurrentSignedPreKeyId()) - bundle['spk_signature'] = signedPreKey.getSignature() - bundle['spk'] = {'key': signedPreKey.getKeyPair().getPublicKey().serialize(), - 'id': signedPreKey.getId()} + bundle['spk_signature'] = spk.getSignature() + bundle['spk'] = {'key': spk.getKeyPair().getPublicKey().serialize(), + 'id': spk.getId()} return OMEMOBundle(**bundle) - def decrypt_msg(self, omemo_message, jid): - own_id = self.own_device_id - if omemo_message.sid == own_id: + def decrypt_message(self, omemo_message, jid): + if omemo_message.sid == self.own_device_id: log.info('Received previously sent message by us') - return - if own_id not in omemo_message.keys: - log.warning('OMEMO message does not contain our device key') - return + raise SelfMessage - encrypted_key, prekey = omemo_message.keys[own_id] + try: + encrypted_key, prekey = omemo_message.keys[self.own_device_id] + except KeyError: + log.info('Received message not for our device') + raise MessageNotForDevice - if prekey: - try: - key = self.handlePreKeyWhisperMessage( + try: + if prekey: + key = self._process_pre_key_message( jid, omemo_message.sid, encrypted_key) - except Exception as error: - log.warning(error) - return - - else: - try: - key = self.handleWhisperMessage( + else: + key = self._process_message( jid, omemo_message.sid, encrypted_key) - except Exception as error: - log.warning(error) - return + + except DuplicateMessageException: + log.info('Received duplicated message') + raise DuplicateMessage + + except Exception as error: + log.warning(error) + raise DecryptionFailed if omemo_message.payload is None: - result = None log.debug("Decrypted Key Exchange Message") - else: - result = aes_decrypt(key, omemo_message.iv, omemo_message.payload) - log.debug("Decrypted Message => %s", result) + raise KeyExchangeMessage + + result = aes_decrypt(key, omemo_message.iv, omemo_message.payload) + log.debug("Decrypted Message => %s", result) return result - def create_msg(self, from_jid, jid, plaintext): + def create_msg(self, jid, plaintext): encrypted_keys = {} devices_list = self.device_list_for(jid) @@ -215,7 +210,7 @@ class OmemoState: for device in devices_list: try: if self.isTrusted(jid, device) == TRUSTED: - cipher = self.get_session_cipher(jid, device) + cipher = self._get_session_cipher(jid, device) cipher_key = cipher.encrypt(result.key) prekey = isinstance(cipher_key, PreKeyWhisperMessage) encrypted_keys[device] = (cipher_key.serialize(), prekey) @@ -233,14 +228,14 @@ class OmemoState: # Encrypt the message key with for each of our own devices for device in my_other_devices: try: - if self.isTrusted(from_jid, device) == TRUSTED: - cipher = self.get_session_cipher(from_jid, device) + if self.isTrusted(self.own_jid, device) == TRUSTED: + cipher = self._get_session_cipher(self.own_jid, device) cipher_key = cipher.encrypt(result.key) prekey = isinstance(cipher_key, PreKeyWhisperMessage) encrypted_keys[device] = (cipher_key.serialize(), prekey) else: log.debug('Skipped own Device because Trust is: %s', - self.isTrusted(from_jid, device)) + self.isTrusted(self.own_jid, device)) except Exception: log.warning('Failed to find key for device: %s', device) @@ -260,7 +255,7 @@ class OmemoState: result = aes_encrypt(plaintext) for tup in devices_list: - self.get_session_cipher(tup[0], tup[1]) + self._get_session_cipher(tup[0], tup[1]) # Encrypt the message key with for each of receivers devices for nick in self.xmpp_con.groupchat[room]: @@ -269,9 +264,9 @@ class OmemoState: continue if jid_to in encrypted_jids: # We already encrypted to this JID continue - if jid_to not in self.session_ciphers: + if jid_to not in self._session_ciphers: continue - for rid, cipher in self.session_ciphers[jid_to].items(): + for rid, cipher in self._session_ciphers[jid_to].items(): try: if self.isTrusted(jid_to, rid) == TRUSTED: cipher_key = cipher.encrypt(result.key) @@ -289,7 +284,7 @@ class OmemoState: # Encrypt the message key with for each of our own devices for dev in my_other_devices: try: - cipher = self.get_session_cipher(from_jid, dev) + cipher = self._get_session_cipher(from_jid, dev) if self.isTrusted(from_jid, dev) == TRUSTED: cipher_key = cipher.encrypt(result.key) prekey = isinstance(cipher_key, PreKeyWhisperMessage) @@ -382,23 +377,21 @@ class OmemoState: self.account, jid, missing_devices) return missing_devices - def get_session_cipher(self, jid, device_id): - if jid not in self.session_ciphers: - self.session_ciphers[jid] = {} - - if device_id not in self.session_ciphers[jid]: + def _get_session_cipher(self, jid, device_id): + try: + return self._session_ciphers[jid][device_id] + except KeyError: cipher = SessionCipher(self.store, self.store, self.store, self.store, jid, device_id) - self.session_ciphers[jid][device_id] = cipher + self._session_ciphers[jid][device_id] = cipher + return cipher - return self.session_ciphers[jid][device_id] - - def handlePreKeyWhisperMessage(self, recipient_id, device_id, key): + def _process_pre_key_message(self, recipient_id, device_id, key): preKeyWhisperMessage = PreKeyWhisperMessage(serialized=key) if not preKeyWhisperMessage.getPreKeyId(): raise Exception('Received PreKeyWhisperMessage ' 'without PreKey => %s' % recipient_id) - sessionCipher = self.get_session_cipher(recipient_id, device_id) + sessionCipher = self._get_session_cipher(recipient_id, device_id) try: log.debug('%s => Received PreKeyWhisperMessage from %s', self.account, recipient_id) @@ -413,12 +406,12 @@ class OmemoState: 'from Untrusted Fingerprint! => %s', self.account, error.getName()) - def handleWhisperMessage(self, recipient_id, device_id, key): + def _process_message(self, recipient_id, device_id, key): whisperMessage = WhisperMessage(serialized=key) log.debug('%s => Received WhisperMessage from %s', self.account, recipient_id) if self.isTrusted(recipient_id, device_id): - sessionCipher = self.get_session_cipher(recipient_id, device_id) + sessionCipher = self._get_session_cipher(recipient_id, device_id) key = sessionCipher.decryptMsg(whisperMessage, textMsg=False) self.add_device(recipient_id, device_id) return key @@ -426,24 +419,24 @@ class OmemoState: raise Exception('Received WhisperMessage ' 'from Untrusted Fingerprint! => %s' % recipient_id) - def checkPreKeyAmount(self): + def _check_pre_key_count(self): # Check if enough PreKeys are available - preKeyCount = self.store.getPreKeyCount() - if preKeyCount < MIN_PREKEY_AMOUNT: - newKeys = DEFAULT_PREKEY_AMOUNT - preKeyCount - self.store.generateNewPreKeys(newKeys) - log.info('%s => %s PreKeys created', self.account, newKeys) + pre_key_count = self.store.getPreKeyCount() + if pre_key_count < MIN_PREKEY_AMOUNT: + missing_count = DEFAULT_PREKEY_AMOUNT - pre_key_count + self.store.generateNewPreKeys(missing_count) + log.info('%s => %s PreKeys created', self.account, missing_count) - def cycleSignedPreKey(self, identityKeyPair): + def _cycle_signed_pre_key(self, ik_pair): # Publish every SPK_CYCLE_TIME a new SignedPreKey # Delete all exsiting SignedPreKeys that are older # then SPK_ARCHIVE_TIME # Check if SignedPreKey exist and create if not if not self.store.getCurrentSignedPreKeyId(): - signedPreKey = KeyHelper.generateSignedPreKey( - identityKeyPair, self.store.getNextSignedPreKeyId()) - self.store.storeSignedPreKey(signedPreKey.getId(), signedPreKey) + spk = KeyHelper.generateSignedPreKey( + ik_pair, self.store.getNextSignedPreKeyId()) + self.store.storeSignedPreKey(spk.getId(), spk) log.debug('%s => New SignedPreKey created, because none existed', self.account) @@ -453,9 +446,9 @@ class OmemoState: self.store.getCurrentSignedPreKeyId()) if int(timestamp) < now - SPK_CYCLE_TIME: - signedPreKey = KeyHelper.generateSignedPreKey( - identityKeyPair, self.store.getNextSignedPreKeyId()) - self.store.storeSignedPreKey(signedPreKey.getId(), signedPreKey) + spk = KeyHelper.generateSignedPreKey( + ik_pair, self.store.getNextSignedPreKeyId()) + self.store.storeSignedPreKey(spk.getId(), spk) log.debug('%s => Cycled SignedPreKey', self.account) # Delete all SignedPreKeys that are older than SPK_ARCHIVE_TIME @@ -465,3 +458,27 @@ class OmemoState: class NoValidSessions(Exception): pass + + +class SelfMessage(Exception): + pass + + +class MessageNotForDevice(Exception): + pass + + +class DecryptionFailed(Exception): + pass + + +class KeyExchangeMessage(Exception): + pass + + +class InvalidMessage(Exception): + pass + + +class DuplicateMessage(Exception): + pass diff --git a/omemo/modules/omemo.py b/omemo/modules/omemo.py index b8ab57f..8254a88 100644 --- a/omemo/modules/omemo.py +++ b/omemo/modules/omemo.py @@ -39,8 +39,14 @@ from gajim.common.modules.base import BaseModule from gajim.common.modules.util import event_node from omemo.backend.state import OmemoState +from omemo.backend.state import KeyExchangeMessage +from omemo.backend.state import SelfMessage +from omemo.backend.state import MessageNotForDevice +from omemo.backend.state import DecryptionFailed +from omemo.backend.state import DuplicateMessage from omemo.modules.util import prepare_stanza + ALLOWED_TAGS = [ ('request', nbxmpp.NS_RECEIPTS), ('active', nbxmpp.NS_CHATSTATES), @@ -159,65 +165,73 @@ class OMEMO(BaseModule): return if properties.is_mam_message: - if properties.omemo.sid == self.omemo.own_device_id: - log.info('%s => Skip message because it was sent by us', - self._account) - raise NodeProcessed - log.info('%s => Message received, archive: %s', - self._account, properties.mam.archive) + from_jid = self._process_mam_message(properties) + elif properties.from_muc: + from_jid = self._process_muc_message(properties) else: - log.info('%s => Message received', self._account) + from_jid = properties.jid.getBare() - from_jid = properties.jid.getBare() + if from_jid is None: + return - if properties.from_muc: - if properties.is_mam_message: - log.info('%s => MUC MAM Message received', self._account) - if properties.muc_user.jid is None: - log.info('No real jid found, ignore message') - return - from_jid = properties.muc_user.jid.getBare() - else: - room_jid = properties.jid.getBare() - resource = properties.jid.getResource() - if properties.muc_ofrom is not None: - # History Message from MUC - from_jid = properties.muc_ofrom.getBare() - else: - try: - from_jid = self.groupchat[room_jid][resource] - except KeyError: - log.debug('Groupchat: Last resort trying to ' - 'find SID in DB') - from_jid = self.omemo.store.getJidFromDevice( - properties.omemo.sid) - if not from_jid: - log.error("%s => Can't decrypt GroupChat Message " - "from %s", self._account, resource) - return - self.groupchat[room_jid][resource] = from_jid + log.info('%s => Message received from: %s', self._account, from_jid) - log.debug('GroupChat Message from: %s', from_jid) + try: + return self.omemo.decrypt_message(properties.omemo, + from_jid) + except (KeyExchangeMessage, DuplicateMessage): + raise NodeProcessed - plaintext = '' - if properties.omemo.sid == self.omemo.own_device_id: - if properties.omemo.payload in self.gc_message: - plaintext = self.gc_message[properties.omemo.payload] - del self.gc_message[properties.omemo.payload] - else: - log.error("%s => Can't decrypt own GroupChat Message", - self._account) - return - else: - plaintext = self.omemo.decrypt_msg(properties.omemo, from_jid) + except SelfMessage: + if properties.from_muc: + if properties.omemo.payload in self.gc_message: + plaintext = self.gc_message[properties.omemo.payload] + del self.gc_message[properties.omemo.payload] + return plaintext - if not plaintext: + log.warning("%s => Can't decrypt own GroupChat Message", + self._account) + raise NodeProcessed + + except (DecryptionFailed, MessageNotForDevice): return prepare_stanza(stanza, plaintext) - self.print_msg_to_log(stanza) + self._debug_print_stanza(stanza) properties.encrypted = EncryptionData({'name': ENCRYPTION_NAME}) + def _process_muc_message(self, properties): + room_jid = properties.jid.getBare() + resource = properties.jid.getResource() + if properties.muc_ofrom is not None: + # History Message from MUC + return properties.muc_ofrom.getBare() + + try: + return self.groupchat[room_jid][resource] + except KeyError: + log.info('%s => Groupchat: Last resort trying to ' + 'find SID in DB', self._account) + from_jid = self.omemo.store.getJidFromDevice(properties.omemo.sid) + if not from_jid: + log.error("%s => Can't decrypt GroupChat Message " + "from %s", self._account, resource) + return + self.groupchat[room_jid][resource] = from_jid + return from_jid + + def _process_mam_message(self, properties): + log.info('%s => Message received, archive: %s', + self._account, properties.mam.archive) + from_jid = properties.jid.getBare() + if properties.from_muc: + log.info('%s => MUC MAM Message received', self._account) + if properties.muc_user.jid is None: + log.info('%s => No real jid found', self._account) + return + from_jid = properties.muc_user.jid.getBare() + return from_jid + def _on_muc_user_presence(self, _con, _stanza, properties): if properties.type == PresenceType.ERROR: return @@ -371,7 +385,7 @@ class OMEMO(BaseModule): create_omemo_message(event.msg_iq, omemo_message, node_whitelist=ALLOWED_TAGS) - self.print_msg_to_log(event.msg_iq) + self._debug_print_stanza(event.msg_iq) callback(event) def encrypt_message(self, conn, event, callback): @@ -385,8 +399,7 @@ class OMEMO(BaseModule): to_jid = app.get_jid_without_resource(event.jid) try: - omemo_message = self.omemo.create_msg( - self.own_jid, to_jid, event.message) + omemo_message = self.omemo.create_msg(to_jid, event.message) if omemo_message is None: raise OMEMOError('Error while encrypting') @@ -401,7 +414,7 @@ class OMEMO(BaseModule): create_omemo_message(event.msg_iq, omemo_message, node_whitelist=ALLOWED_TAGS) - self.print_msg_to_log(event.msg_iq) + self._debug_print_stanza(event.msg_iq) event.xhtml = None event.encrypted = ENCRYPTION_NAME event.additional_data['encrypted'] = {'name': ENCRYPTION_NAME} @@ -561,7 +574,7 @@ class OMEMO(BaseModule): self.are_keys_missing(jid) @staticmethod - def print_msg_to_log(stanza): + def _debug_print_stanza(stanza): log.debug('-'*15) stanzastr = '\n' + stanza.__str__(fancy=True) stanzastr = stanzastr[0:-1]