[omemo] More refactoring

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
-
-        return self.session_ciphers[jid][device_id]
+            self._session_ciphers[jid][device_id] = cipher
+            return cipher
 
-    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

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.debug('GroupChat Message from: %s', from_jid)
-
-        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)
+        log.info('%s => Message received from: %s', self._account, from_jid)
 
-        if not plaintext:
+        try:
+            return self.omemo.decrypt_message(properties.omemo,
+                                              from_jid)
+        except (KeyExchangeMessage, DuplicateMessage):
+            raise NodeProcessed
+
+        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
+
+                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]