mam.py 6.86 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 29 30 31 32 33 34 35
def add_line(
        tab,
        text_buffer: TextBuffer,
        text: str,
        time: datetime,
        nick: str,
        top: bool,
    ) -> None:
36 37
    """Adds a textual entry in the TextBuffer"""

38
    # Convert to local timezone
39 40
    time = time.replace(tzinfo=timezone.utc).astimezone(tz=None)
    time = time.replace(tzinfo=None)
41

Maxime Buquet's avatar
Maxime Buquet committed
42
    deterministic = config.get_by_tabname('deterministic_nick_colors', tab.jid.bare)
43 44 45
    if isinstance(tab, tabs.MucTab):
        nick = nick.split('/')[1]
        user = tab.get_user_by_name(nick)
46 47 48 49 50 51 52 53 54 55
        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
56
                    nick_pos = int(md5(nick.encode('utf-8')).hexdigest(), 16) % mod
57
                    color = theme.LIST_COLOR_NICKNAMES[nick_pos]
58 59 60 61 62 63 64
        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
65
    text_buffer.add_message(
66 67 68 69 70 71 72 73 74 75 76
        txt=text,
        time=time,
        nickname=nick,
        nick_color=color,
        history=True,
        user=None,
        highlight=False,
        top=top,
        identifier=None,
        str_time=None,
        jid=None,
77 78
    )

Maxime Buquet's avatar
Maxime Buquet committed
79

Madhur Garg's avatar
Madhur Garg committed
80 81 82 83 84 85 86 87 88 89 90
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:
91
    try:
Madhur Garg's avatar
Madhur Garg committed
92
        iq = await core.xmpp.plugin['xep_0030'].get_info(jid=remote_jid)
93
    except (IqError, IqTimeout):
Madhur Garg's avatar
Madhur Garg committed
94 95 96 97 98 99 100 101 102 103 104
        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
105
    else:
Madhur Garg's avatar
Madhur Garg committed
106 107 108 109 110 111
        args['with_jid'] = remote_jid

    args['rsm'] = {'max': amount}
    if reverse:
        if before is not None:
            args['rsm']['before'] = before
112
        else:
Madhur Garg's avatar
Madhur Garg committed
113 114 115 116 117 118 119 120 121 122 123 124 125 126
            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
127

128
async def add_messages_to_buffer(tab, top: bool, results, amount: int) -> bool:
Madhur Garg's avatar
Madhur Garg committed
129 130 131
    """Prepends or appends messages to the tab text_buffer"""

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

Maxime Buquet's avatar
Maxime Buquet committed
162

Madhur Garg's avatar
Madhur Garg committed
163
async def fetch_history(tab, end: Optional[datetime] = None, amount: Optional[int] = None):
164
    remote_jid = tab.jid
165
    before = tab.last_stanza_id
Madhur Garg's avatar
Madhur Garg committed
166 167
    if end is None:
        end = datetime.now()
168 169 170 171
    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
172

173 174
    if amount >= 100:
        amount = 99
Madhur Garg's avatar
Madhur Garg committed
175 176 177

    groupchat = isinstance(tab, tabs.MucTab)

Maxime Buquet's avatar
Maxime Buquet committed
178 179 180 181 182 183 184 185 186
    results = await query(
        tab.core,
        groupchat,
        remote_jid,
        amount,
        reverse=True,
        end=end,
        before=before,
    )
Madhur Garg's avatar
Madhur Garg committed
187 188 189
    query_status = await add_messages_to_buffer(tab, True, results, amount)
    tab.query_status = query_status

Maxime Buquet's avatar
Maxime Buquet committed
190

Madhur Garg's avatar
Madhur Garg committed
191 192 193 194 195 196 197 198 199 200 201 202 203
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
204

Madhur Garg's avatar
Madhur Garg committed
205
async def on_scroll_up(tab) -> None:
206 207
    tw = tab.text_win

208 209 210
    # 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.
211 212 213 214 215 216
    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
217
    try:
218 219 220 221
        # 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
222 223 224 225 226 227
    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