Improve XEP-0308 support

- Prevent correction of delayed messages
- Prevent correction of messages by someone else in a MUC (and in a
  private tab)
- Messages with unauthorized corrections (above) or wrong message id
  will be displayed as normal messages

TODO: restrict the corrections to the same fullJID (only in direct
"normal" conversations, because we can know in private an muc tabs, via
the User object)

3 files changed, 60 insertions, 50 deletions

diff --git a/src/core.py b/src/core.py
index c7fac492..4d1bf45a 100644
--- a/src/core.py
+++ b/src/core.py
@@ -53,7 +53,7 @@ from config import config
 from logger import logger
 from roster import roster
 from contact import Contact, Resource
-from text_buffer import TextBuffer
+from text_buffer import TextBuffer, CorrectionError
 from keyboard import keyboard
 from theming import get_theme
 from fifo import Fifo
@@ -2551,9 +2551,14 @@ class Core(object):
             delayed = False
             date = None
         replaced_id = message['replace']['id']
+        replaced = False
         if replaced_id is not '':
-            conversation.modify_message(body, replaced_id, message['id'])
-        else:
+            try:
+                conversation.modify_message(body, replaced_id, message['id'])
+                replaced = True
+            except CorrectionError:
+                pass
+        if not replaced :
             conversation._text_buffer.add_message(body, date,
                     nickname=remote_nick,
                     nick_color=get_theme().COLOR_REMOTE_USER,
@@ -2603,18 +2608,23 @@ class Core(object):
         if not tab:
             self.information(_("message received for a non-existing room: %s") % (room_from))
             return
-        if tab.get_user_by_name(nick_from) and\
-                tab.get_user_by_name(nick_from) in tab.ignores:
+        user = tab.get_user_by_name(nick_from)
+        if user and user in tab.ignores:
             return
         self.events.trigger('muc_msg', message, tab)
         body = xhtml.get_body_from_message_stanza(message)
         if body:
             date = date if delayed == True else None
             replaced_id = message['replace']['id']
+            replaced = False
             if replaced_id is not '':
-                if tab.modify_message(body, replaced_id, message['id'], date, nick_from):
-                    self.events.trigger('highlight', message, tab)
-            elif tab.add_message(body, date, nick_from, history=True if date else False, identifier=message['id']):
+                try:
+                    if tab.modify_message(body, replaced_id, message['id'], date, nick_from, user):
+                        self.events.trigger('highlight', message, tab)
+                    replaced = True
+                except CorrectionError:
+                    pass
+            if not replaced and tab.add_message(body, date, nick_from, history=True if date else False, identifier=message['id']):
                 self.events.trigger('highlight', message, tab)
             if tab is self.current_tab():
                 tab.text_win.refresh()
@@ -2653,26 +2663,31 @@ class Core(object):
         if not body or not tab:
             return
         replaced_id = message['replace']['id']
+        replaced = False
         if replaced_id is not '':
-            tab.modify_message(body, replaced_id, message['id'])
-        else:
+            user = self.get_tab_by_name(room_from, tabs.MucTab).get_user_by_name(nick_from),
+            try:
+                tab.modify_message(body, replaced_id, message['id'], user=user)
+                replaced = True
+            except CorrectionError:
+                pass
+        if not replaced:
             tab.add_message(body, time=None, nickname=nick_from,
                             forced_user=self.get_tab_by_name(room_from, tabs.MucTab).get_user_by_name(nick_from),
                             identifier=message['id'])
-        conversation = self.get_tab_by_name(jid.full, tabs.PrivateTab)
-        if conversation and conversation.remote_wants_chatstates is None:
+        if tab.remote_wants_chatstates is None:
             if message['chat_state']:
-                conversation.remote_wants_chatstates = True
+                tab.remote_wants_chatstates = True
             else:
-                conversation.remote_wants_chatstates = False
+                tab.remote_wants_chatstates = False
         if 'private' in config.get('beep_on', 'highlight private').split():
             if config.get_by_tabname('disable_beep', 'false', jid.full, False).lower() != 'true':
                 curses.beep()
         logger.log_message(jid.full.replace('/', '\\'), nick_from, body)
-        if conversation is self.current_tab():
+        if tab is self.current_tab():
             self.refresh_window()
         else:
-            conversation.state = 'private'
+            tab.state = 'private'
             self.refresh_tab_win()
 
     ### Chatstates ###
diff --git a/src/tabs.py b/src/tabs.py
index b81bce8b..db840911 100644
--- a/src/tabs.py
+++ b/src/tabs.py
@@ -482,6 +482,23 @@ class ChatTab(Tab):
         if time is None and self.joined:        # don't log the history messages
             logger.log_message(self.name, nickname, txt)
 
+    def add_message(self, txt, time=None, nickname=None, forced_user=None, nick_color=None, identifier=None):
+        self._text_buffer.add_message(txt, time=time,
+                nickname=nickname,
+                nick_color=nick_color,
+                history=None,
+                user=forced_user,
+                identifier=identifier)
+
+    def modify_message(self, txt, old_id, new_id, user=None):
+        self.log_message(txt, time, self.name)
+        message = self._text_buffer.modify_message(txt, old_id, new_id, time=time, user=user)
+        if message:
+            self.text_win.modify_message(old_id, message)
+            self.core.refresh_window()
+            return True
+        return False
+
     def last_words_completion(self):
         """
         Complete the input with words recently said
@@ -1733,14 +1750,15 @@ class MucTab(ChatTab):
         self._text_buffer.add_message(txt, time, nickname, nick_color, history, user, highlight=highlight, identifier=identifier)
         return highlight
 
-    def modify_message(self, txt, old_id, new_id, time=None, nickname=None):
+    def modify_message(self, txt, old_id, new_id, time=None, nickname=None, user=None):
         self.log_message(txt, time, nickname)
         highlight = self.do_highlight(txt, time, nickname)
-        message = self._text_buffer.modify_message(txt, old_id, new_id, highlight=highlight, time=time)
+        message = self._text_buffer.modify_message(txt, old_id, new_id, highlight=highlight, time=time, user=user)
         if message:
             self.text_win.modify_message(old_id, message)
             self.core.refresh_window()
             return highlight
+        return False
 
     def matching_names(self):
         return [safeJID(self.get_name()).user]
@@ -1826,6 +1844,7 @@ class PrivateTab(ChatTab):
         else:
             self.add_message(msg['body'],
                     nickname=self.core.own_nick or self.own_nick,
+                    forced_user=self.parent_muc.get_user_by_name(self.own_nick),
                     nick_color=get_theme().COLOR_OWN_NICK,
                     identifier=msg['id'])
         if msg['body'].find('\x19') != -1:
@@ -2013,21 +2032,6 @@ class PrivateTab(ChatTab):
         if reason:
             self.add_message(txt=reason)
 
-    def modify_message(self, txt, old_id, new_id):
-        self.log_message(txt, time, self.name)
-        message = self._text_buffer.modify_message(txt, old_id, new_id, time=time)
-        if message:
-            self.text_win.modify_message(old_id, message)
-            self.core.refresh_window()
-
-    def add_message(self, txt, time=None, nickname=None, forced_user=None, nick_color=None, identifier=None):
-        self._text_buffer.add_message(txt, time=time,
-                nickname=nickname,
-                nick_color=nick_color,
-                history=None,
-                user=forced_user,
-                identifier=identifier)
-
     def matching_names(self):
         return [safeJID(self.get_name()).resource]
 
@@ -3175,21 +3179,6 @@ class ConversationTab(ChatTab):
         if config.get_by_tabname('send_chat_states', 'true', self.general_jid, True) == 'true':
             self.send_chat_state('gone')
 
-    def modify_message(self, txt, old_id, new_id):
-        self.log_message(txt, time, self.name)
-        message = self._text_buffer.modify_message(txt, old_id, new_id, time=time)
-        if message:
-            self.text_win.modify_message(old_id, message)
-            self.core.refresh_window()
-
-    def add_message(self, txt, time=None, nickname=None, forced_user=None, nick_color=None, identifier=None):
-        self._text_buffer.add_message(txt, time=time,
-                nickname=nickname,
-                nick_color=nick_color,
-                history=None,
-                user=forced_user,
-                identifier=identifier)
-
     def matching_names(self):
         contact = roster[self.get_name()]
         res = [contact.bare_jid if contact else safeJID(self.get_name()).bare]
diff --git a/src/text_buffer.py b/src/text_buffer.py
index 2aa80670..0b9fc319 100644
--- a/src/text_buffer.py
+++ b/src/text_buffer.py
@@ -21,6 +21,8 @@ from theming import get_theme
 message_fields = 'txt nick_color time str_time nickname user identifier highlight me old_message revisions'
 Message = collections.namedtuple('Message', message_fields)
 
+class CorrectionError(Exception): pass
+
 def other_elems(self):
     acc = ['Message(']
     fields = message_fields.split()
@@ -103,16 +105,20 @@ class TextBuffer(object):
                 window.scroll_up(nb)
         return ret_val or 1
 
-    def modify_message(self, txt, old_id, new_id, highlight=False, time=None):
+    def modify_message(self, txt, old_id, new_id, highlight=False, time=None, user=None):
         for i in range(len(self.messages) -1, -1, -1):
             msg = self.messages[i]
             if msg.identifier == old_id:
+                if msg.user and msg.user is not user:
+                    raise CorrectionError("wrong user")
+                elif len(msg.str_time) > 8: # ugly
+                    raise CorrectionError("delayed message")
                 message = self.make_message(txt, time if time else msg.time, msg.nickname, msg.nick_color, None, msg.user, new_id, highlight=highlight, old_message=msg, revisions=msg.revisions + 1)
                 self.messages[i] = message
                 log.debug('Replacing message %s with %s.', old_id, new_id)
                 return message
         log.debug('Message %s not found in text_buffer, abort replacement.', old_id)
-        return
+        raise CorrectionError("nothing to replace")
 
     def del_window(self, win):
         self.windows.remove(win)