mam.py 6.39 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)
35 36
    deterministic = config.get_by_tabname('deterministic_nick_colors',
                                              tab.jid.bare)
37 38 39
    if isinstance(tab, tabs.MucTab):
        nick = nick.split('/')[1]
        user = tab.get_user_by_name(nick)
40 41 42 43 44 45 46 47 48 49 50 51 52
        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]
53 54 55 56 57 58 59
        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
60
    text_buffer.add_message(
61 62 63 64 65 66 67 68 69 70 71
        txt=text,
        time=time,
        nickname=nick,
        nick_color=color,
        history=True,
        user=None,
        highlight=False,
        top=top,
        identifier=None,
        str_time=None,
        jid=None,
72 73
    )

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

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

async def add_messages_to_buffer(tab, top: bool, results, amount: int) -> None:
    """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

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

165 166
    if amount >= 100:
        amount = 99
Madhur Garg's avatar
Madhur Garg committed
167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196

    groupchat = isinstance(tab, tabs.MucTab)

    results = await query(tab.core, groupchat, remote_jid, amount, reverse=True, end=end, before=before)
    query_status = await add_messages_to_buffer(tab, True, results, amount)
    tab.query_status = query_status

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

async def on_scroll_up(tab) -> None:
    amount = tab.text_win.height
    try:
        await fetch_history(tab, amount=amount)
    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