mam.py 6.88 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 random
10
from datetime import datetime, timedelta, timezone
11 12 13
from hashlib import md5
from typing import Optional, Callable

Madhur Garg's avatar
Madhur Garg committed
14
from slixmpp import JID
15
from slixmpp.exceptions import IqError, IqTimeout
16
from poezio.theming import get_theme
17
from poezio import tabs
18
from poezio import xhtml, colors
19
from poezio.config import config
20
from poezio.text_buffer import TextBuffer
Madhur Garg's avatar
Madhur Garg committed
21 22 23 24 25 26


class DiscoInfoException(Exception): pass
class MAMQueryException(Exception): pass
class NoMAMSupportException(Exception): pass

27

28
def add_line(tab, text_buffer: TextBuffer, text: str, str_time: str, nick: str, top: bool):
29 30 31 32
    """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')
33 34
    time = time.replace(tzinfo=timezone.utc).astimezone(tz=None)
    time = time.replace(tzinfo=None)
Maxime Buquet's avatar
Maxime Buquet committed
35
    deterministic = config.get_by_tabname('deterministic_nick_colors', tab.jid.bare)
36 37 38
    if isinstance(tab, tabs.MucTab):
        nick = nick.split('/')[1]
        user = tab.get_user_by_name(nick)
39 40 41 42 43 44 45 46 47 48
        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)
Maxime Buquet's avatar
Maxime Buquet committed
49
                    nick_pos = int(md5(nick.encode('utf-8')).hexdigest(), 16) % mod
50
                    color = theme.LIST_COLOR_NICKNAMES[nick_pos]
51 52 53 54 55 56 57
        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
58
    text_buffer.add_message(
59 60 61 62 63 64 65 66 67 68 69
        txt=text,
        time=time,
        nickname=nick,
        nick_color=color,
        history=True,
        user=None,
        highlight=False,
        top=top,
        identifier=None,
        str_time=None,
        jid=None,
70 71
    )

Maxime Buquet's avatar
Maxime Buquet committed
72

Madhur Garg's avatar
Madhur Garg committed
73 74 75 76 77 78 79 80 81 82 83
async def query(
        core,
        groupchat: bool,
        remote_jid: JID,
        amount: int,
        reverse: bool,
        start: Optional[datetime] = None,
        end: Optional[datetime] = None,
        before: Optional[str] = None,
        callback: Optional[Callable] = None,
    ) -> None:
84
    try:
Madhur Garg's avatar
Madhur Garg committed
85
        iq = await core.xmpp.plugin['xep_0030'].get_info(jid=remote_jid)
86
    except (IqError, IqTimeout):
Madhur Garg's avatar
Madhur Garg committed
87 88 89 90 91 92 93 94 95 96 97
        raise DiscoInfoException
    if 'urn:xmpp:mam:2' not in iq['disco_info'].get_features():
        raise NoMAMSupportException

    args = {
        'iterator': True,
        'reverse': reverse,
    }

    if groupchat:
        args['jid'] = remote_jid
98
    else:
Madhur Garg's avatar
Madhur Garg committed
99 100 101 102 103 104
        args['with_jid'] = remote_jid

    args['rsm'] = {'max': amount}
    if reverse:
        if before is not None:
            args['rsm']['before'] = before
105
        else:
Madhur Garg's avatar
Madhur Garg committed
106 107 108 109 110 111 112 113 114 115 116 117 118 119
            args['end'] = end
    else:
        args['rsm']['start'] = start
        if before is not None:
            args['rsm']['end'] = end
    try:
        results = core.xmpp['xep_0313'].retrieve(**args)
    except (IqError, IqTimeout):
        raise MAMQueryException
    if callback is not None:
        callback(results)

    return results

Maxime Buquet's avatar
Maxime Buquet committed
120

121
async def add_messages_to_buffer(tab, top: bool, results, amount: int) -> bool:
Madhur Garg's avatar
Madhur Garg committed
122 123 124
    """Prepends or appends messages to the tab text_buffer"""

    text_buffer = tab._text_buffer
125 126
    msg_count = 0
    msgs = []
127
    async for rsm in results:
128 129
        if top:
            for msg in rsm['mam']['results']:
Madhur Garg's avatar
Madhur Garg committed
130 131
                if msg['mam_result']['forwarded']['stanza'] \
                .xml.find('{%s}%s' % ('jabber:client', 'body')) is not None:
132
                    msgs.append(msg)
133
                if msg_count == amount:
134
                    tab.core.refresh_window()
Madhur Garg's avatar
Madhur Garg committed
135
                    return False
136 137 138 139 140 141
                msg_count += 1
            msgs.reverse()
            for msg in msgs:
                forwarded = msg['mam_result']['forwarded']
                timestamp = forwarded['delay']['stamp']
                message = forwarded['stanza']
142
                tab.last_stanza_id = msg['mam_result']['id']
143
                nick = str(message['from'])
144
                add_line(tab, text_buffer, message['body'], timestamp, nick, top)
145 146 147 148 149
        else:
            for msg in rsm['mam']['results']:
                forwarded = msg['mam_result']['forwarded']
                timestamp = forwarded['delay']['stamp']
                message = forwarded['stanza']
150
                nick = str(message['from'])
151
                add_line(tab, text_buffer, message['body'], timestamp, nick, top)
152
                tab.core.refresh_window()
Madhur Garg's avatar
Madhur Garg committed
153
    return False
154

Maxime Buquet's avatar
Maxime Buquet committed
155

Madhur Garg's avatar
Madhur Garg committed
156
async def fetch_history(tab, end: Optional[datetime] = None, amount: Optional[int] = None):
157
    remote_jid = tab.jid
158
    before = tab.last_stanza_id
Madhur Garg's avatar
Madhur Garg committed
159 160
    if end is None:
        end = datetime.now()
161 162 163 164
    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')
Madhur Garg's avatar
Madhur Garg committed
165

166 167
    if amount >= 100:
        amount = 99
Madhur Garg's avatar
Madhur Garg committed
168 169 170

    groupchat = isinstance(tab, tabs.MucTab)

Maxime Buquet's avatar
Maxime Buquet committed
171 172 173 174 175 176 177 178 179
    results = await query(
        tab.core,
        groupchat,
        remote_jid,
        amount,
        reverse=True,
        end=end,
        before=before,
    )
Madhur Garg's avatar
Madhur Garg committed
180 181 182
    query_status = await add_messages_to_buffer(tab, True, results, amount)
    tab.query_status = query_status

Maxime Buquet's avatar
Maxime Buquet committed
183

Madhur Garg's avatar
Madhur Garg committed
184 185 186 187 188 189 190 191 192 193 194 195 196
async def on_tab_open(tab) -> None:
    amount = 2 * tab.text_win.height
    end = datetime.now()
    for message in tab._text_buffer.messages:
        time = message.time
        if time < end:
            end = time
    end = end + timedelta(seconds=-1)
    try:
        await fetch_history(tab, end=end, amount=amount)
    except (NoMAMSupportException, MAMQueryException, DiscoInfoException):
        return None

Maxime Buquet's avatar
Maxime Buquet committed
197

Madhur Garg's avatar
Madhur Garg committed
198
async def on_scroll_up(tab) -> None:
199 200
    tw = tab.text_win

201 202 203
    # If position in the tab is < two screen pages, then fetch MAM, so that we
    # keep some prefetched margin. A first page should also be prefetched on
    # join if not already available.
204 205 206 207 208 209
    total, pos, height = len(tw.built_lines), tw.pos, tw.height
    rest = (total - pos) // height

    if rest > 1:
        return None

Madhur Garg's avatar
Madhur Garg committed
210
    try:
211 212 213 214
        # XXX: Do we want to fetch a possibly variable number of messages?
        # (InfoTab changes height depending on the type of messages, see
        # `information_buffer_popup_on`).
        await fetch_history(tab, amount=height)
Madhur Garg's avatar
Madhur Garg committed
215 216 217 218 219 220
    except NoMAMSupportException:
        tab.core.information('MAM not supported for %r' % tab.jid, 'Info')
        return None
    except (MAMQueryException, DiscoInfoException):
        tab.core.information('An error occured when fetching MAM for %r' % tab.jid, 'Error')
        return None