[omemo] Refactor encryption methods

omemo/omemoplugin.py

@@ -27,6 +27,7 @@ import shutil
 import nbxmpp
 import binascii
 import threading
+import time
 
 from gi.repository import GLib
 from nbxmpp.simplexml import Node
@@ -45,7 +46,7 @@ from .xmpp import (
     unpack_device_list_update, unpack_encrypted)
 
 from common.connection_handlers_events import (
-    MessageReceivedEvent, MamMessageReceivedEvent)
+    MessageReceivedEvent, MamMessageReceivedEvent, MessageNotSentEvent)
 
 
 IQ_CALLBACK = {}
@@ -614,32 +615,39 @@ class OmemoPlugin(GajimPlugin):
             exception or error occurs
         """
         account = event.conn.name
-        if account in self.disabled_accounts:
-            return
         try:
-            if not event.msg_iq.getTag('body'):
+            if account in self.disabled_accounts:
-                return
+                raise OMEMOError('Account disabled in OMEMO config')
-            state = self.get_omemo_state(account)
-            full_jid = str(event.msg_iq.getAttr('to'))
-            to_jid = gajim.get_jid_without_resource(full_jid)
-
-            plaintext = event.msg_iq.getBody()
-            msg_dict = state.create_gc_msg(
-                gajim.get_jid_from_account(account),
-                to_jid,
-                plaintext.encode('utf8'))
-            if not msg_dict:
-                return True
         
             self.cleanup_stanza(event)
 
-            self.gc_message[msg_dict['payload']] = plaintext
+            if not event.message:
+                callback(event)
+                return
+
+            state = self.get_omemo_state(account)
+            to_jid = gajim.get_jid_without_resource(event.jid)
+            own_jid = gajim.get_jid_from_account(account)
+
+            msg_dict = state.create_gc_msg(
+                own_jid, to_jid, event.message.encode('utf8'))
+            if not msg_dict:
+                raise OMEMOError('Error while encrypting')
+
+        except OMEMOError as error:
+            log.error(error)
+            gajim.nec.push_incoming_event(
+                MessageNotSentEvent(
+                    None, conn=conn, jid=event.jid, message=event.message,
+                    error=error, time_=time.time(), session=None))
+            return
+
+        self.gc_message[msg_dict['payload']] = event.message
         encrypted_node = OmemoMessage(msg_dict)
 
         event.msg_iq.addChild(node=encrypted_node)
 
         # XEP-0380: Explicit Message Encryption
-            if not event.msg_iq.getTag('encryption', attrs={'xmlns': NS_EME}):
         eme_node = Node('encryption', attrs={'xmlns': NS_EME,
                                              'name': 'OMEMO',
                                              'namespace': NS_OMEMO})
@@ -655,11 +663,7 @@ class OmemoPlugin(GajimPlugin):
         event.msg_iq.addChild(node=store)
 
         self.print_msg_to_log(event.msg_iq)
-
         callback(event)
-        except Exception as e:
-            log.debug(e)
-            return
 
     def _encrypt_message(self, conn, event, callback):
         """ Manipulates the outgoing stanza
@@ -676,30 +680,37 @@ class OmemoPlugin(GajimPlugin):
             exception or error occurs
         """
         account = event.conn.name
-        if account in self.disabled_accounts:
-            return
         try:
-            if not event.msg_iq.getTag('body'):
+            if account in self.disabled_accounts:
+                raise OMEMOError('Account disabled in OMEMO config')
+
+            self.cleanup_stanza(event)
+
+            if not event.message:
+                callback(event)
                 return
 
             state = self.get_omemo_state(account)
-            full_jid = str(event.msg_iq.getAttr('to'))
+            to_jid = gajim.get_jid_without_resource(event.jid)
-            to_jid = gajim.get_jid_without_resource(full_jid)
+            own_jid = gajim.get_jid_from_account(account)
 
-            plaintext = event.msg_iq.getBody().encode('utf8')
+            plaintext = event.message.encode('utf8')
-
+            msg_dict = state.create_msg(own_jid, to_jid, plaintext)
-            msg_dict = state.create_msg(
-                gajim.get_jid_from_account(account), to_jid, plaintext)
             if not msg_dict:
-                return True
+                raise OMEMOError('Error while encrypting')
+
+        except OMEMOError as error:
+            log.error(error)
+            gajim.nec.push_incoming_event(
+                MessageNotSentEvent(
+                    None, conn=conn, jid=event.jid, message=event.message,
+                    error=error, time_=time.time(), session=event.session))
+            return
 
         encrypted_node = OmemoMessage(msg_dict)
-            self.cleanup_stanza(event)
-
         event.msg_iq.addChild(node=encrypted_node)
 
         # XEP-0380: Explicit Message Encryption
-            if not event.msg_iq.getTag('encryption', attrs={'xmlns': NS_EME}):
         eme_node = Node('encryption', attrs={'xmlns': NS_EME,
                                              'name': 'OMEMO',
                                              'namespace': NS_OMEMO})
@@ -712,8 +723,6 @@ class OmemoPlugin(GajimPlugin):
         event.xhtml = None
         event.encrypted = self.encryption_name
         callback(event)
-        except Exception as e:
-            log.debug(e)
 
     @staticmethod
     def cleanup_stanza(obj):
@@ -1118,3 +1127,7 @@ class OmemoPlugin(GajimPlugin):
         """
         state = self.get_omemo_state(account)
         state.encryption.deactivate(jid)
+
+
+class OMEMOError(Exception):
+    pass