mam.py 6.77 KB
Newer Older
1 2 3 4 5 6 7 8
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
    Query and control an archive of messages stored on a server using
    XEP-0313: Message Archive Management(MAM).
"""

9
import asyncio
10
import random
11
from datetime import datetime, timedelta, timezone
12
from slixmpp.exceptions import IqError, IqTimeout
13
from poezio.theming import get_theme
14
from poezio import tabs
15
from poezio import xhtml, colors
16
from poezio.config import config
17 18
from poezio.text_buffer import Message, TextBuffer

19
def add_line(tab, text_buffer: TextBuffer, text: str, str_time: str, nick: str, top: bool):
20 21 22 23
    """Adds a textual entry in the TextBuffer"""

    time = datetime.strftime(str_time, '%Y-%m-%d %H:%M:%S')
    time = datetime.strptime(time, '%Y-%m-%d %H:%M:%S')
24 25
    time = time.replace(tzinfo=timezone.utc).astimezone(tz=None)
    time = time.replace(tzinfo=None)
26 27
    deterministic = config.get_by_tabname('deterministic_nick_colors',
                                              tab.jid.bare)
28 29 30
    if isinstance(tab, tabs.MucTab):
        nick = nick.split('/')[1]
        user = tab.get_user_by_name(nick)
31 32 33 34 35 36 37 38 39 40 41 42 43
        if deterministic:
            if user:
                color = user.color
            else:
                theme = get_theme()
                if theme.ccg_palette:
                    fg_color = colors.ccg_text_to_color(theme.ccg_palette, nick)
                    color = fg_color, -1
                else:
                    mod = len(theme.LIST_COLOR_NICKNAMES)
                    nick_pos = int(md5(nick.encode('utf-8')).hexdigest(),
                                16) % mod
                    color = theme.LIST_COLOR_NICKNAMES[nick_pos]
44 45 46 47 48 49 50
        else:
            color = random.choice(list(xhtml.colors))
            color = xhtml.colors.get(color)
            color = (color, -1)
    else:
        nick = nick.split('/')[0]
        color = get_theme().COLOR_OWN_NICK
51
    text_buffer.add_message(
52 53 54 55 56 57 58 59 60 61 62
        txt=text,
        time=time,
        nickname=nick,
        nick_color=color,
        history=True,
        user=None,
        highlight=False,
        top=top,
        identifier=None,
        str_time=None,
        jid=None,
63 64
    )

65
async def query(tab, remote_jid, action, amount, top, start=None, end=None, before=None):
66
    text_buffer = tab._text_buffer
67
    try:
68
        iq = await tab.core.xmpp.plugin['xep_0030'].get_info(jid=remote_jid)
69
    except (IqError, IqTimeout):
70 71
        if action is 'scroll':
            return tab.core.information('%s : Failed to retrieve messages' % remote_jid, 'Error')
72
    if 'urn:xmpp:mam:2' not in iq['disco_info'].get_features() and action is 'scroll':
73
        return tab.core.information("%s doesn't support MAM." % remote_jid, "Info")
74
    if top:
75
        if isinstance(tab, tabs.MucTab):
76
            try:
77
                if before is not None:
78
                    results = tab.core.xmpp['xep_0313'].retrieve(jid=remote_jid,
79
                    iterator=True, reverse=top, rsm={'before':before, 'max':amount})
80
                else:
81
                    results = tab.core.xmpp['xep_0313'].retrieve(jid=remote_jid,
82
                    iterator=True, reverse=top, end=end, rsm={'max':amount})
83
            except (IqError, IqTimeout):
84
                if action is 'scroll':
85
                    return tab.core.information('%s : Failed to retrieve messages' % remote_jid, 'Error')
86
        else:
87
            try:
88
                if before is not None:
89
                    results = tab.core.xmpp['xep_0313'].retrieve(with_jid=remote_jid,
90
                    iterator=True, reverse=top, rsm={'before':before, 'max':amount})
91
                else:
92
                    results = tab.core.xmpp['xep_0313'].retrieve(with_jid=remote_jid,
93
                    iterator=True, reverse=top, end=end, rsm={'max':amount})
94
            except (IqError, IqTimeout):
95
                if action is 'scroll':
96
                    return tab.core.information('%s : Failed to retrieve messages' % remote_jid, 'Error')
97
    else:
98
        if 'conference' in list(iq['disco_info']['identities'])[0]:
99
            try:
100
                results = tab.core.xmpp['xep_0313'].retrieve(jid=remote_jid,
101
                iterator=True, reverse=top, start=start, end=end)
102
            except (IqError, IqTimeout):
103
                return tab.core.information('%s : Failed to retrieve messages' % remote_jid, 'Error')
104
        else:
105
            try:
106
                results = tab.core.xmpp['xep_0313'].retrieve(with_jid=remote_jid,
107
                iterator=True, reverse=top, start=start, end=end)
108
            except (IqError, IqTimeout):
109
                return tab.core.information('%s : Failed to retrieve messages' % remote_jid, 'Error')
110 111
    msg_count = 0
    msgs = []
112
    async for rsm in results:
113 114
        if top:
            for msg in rsm['mam']['results']:
115 116
                if msg['mam_result']['forwarded']['stanza'].xml.find(
                    '{%s}%s' % ('jabber:client', 'body')) is not None:
117
                    msgs.append(msg)
118
                if msg_count == amount:
119
                    tab.query_status = False
120
                    tab.core.refresh_window()
121 122 123 124 125 126 127
                    return
                msg_count += 1
            msgs.reverse()
            for msg in msgs:
                forwarded = msg['mam_result']['forwarded']
                timestamp = forwarded['delay']['stamp']
                message = forwarded['stanza']
128
                tab.last_stanza_id = msg['mam_result']['id']
129
                nick = str(message['from'])
130
                add_line(tab, text_buffer, message['body'], timestamp, nick, top)
131 132 133 134 135
        else:
            for msg in rsm['mam']['results']:
                forwarded = msg['mam_result']['forwarded']
                timestamp = forwarded['delay']['stamp']
                message = forwarded['stanza']
136
                nick = str(message['from'])
137
                add_line(tab, text_buffer, message['body'], timestamp, nick, top)
138
                tab.core.refresh_window()
139
    tab.query_status = False
140

141
def mam_scroll(tab, action):
142 143
    remote_jid = tab.jid
    text_buffer = tab._text_buffer
144
    before = tab.last_stanza_id
145 146 147 148 149 150
    end = datetime.now()
    if isinstance(tab, tabs.MucTab) is False:
        for message in text_buffer.messages:
            time = message.time
            if time < end:
                end = time
151
        end = end + timedelta(seconds=-1)
152 153 154 155
    tzone = datetime.now().astimezone().tzinfo
    end = end.replace(tzinfo=tzone).astimezone(tz=timezone.utc)
    end = end.replace(tzinfo=None)
    end = datetime.strftime(end, '%Y-%m-%dT%H:%M:%SZ')
156 157 158 159 160 161 162 163 164 165 166
    if action is 'scroll':
        amount = tab.text_win.height
    else:
        amount = 2 * tab.text_win.height
    if amount >= 100:
        amount = 99
    if before is None:
        asyncio.ensure_future(query(tab, remote_jid, action, amount, top=True, end=end))
    else:
        asyncio.ensure_future(query(tab, remote_jid, action, amount, top=True, before=before))
    tab.query_status = True