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 asyncio
10
import random
11
from datetime import datetime, timedelta, timezone
Madhur Garg's avatar
Madhur Garg committed
12
from slixmpp import JID
13
from slixmpp.exceptions import IqError, IqTimeout
14
from poezio.theming import get_theme
15
from poezio import tabs
16
from poezio import xhtml, colors
17
from poezio.config import config
18
from poezio.text_buffer import Message, TextBuffer
Madhur Garg's avatar
Madhur Garg committed
19 20 21 22 23 24 25
from typing import List, Optional, Callable


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

26

27
def add_line(tab, text_buffer: TextBuffer, text: str, str_time: str, nick: str, top: bool):
28 29 30 31
    """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')
32 33
    time = time.replace(tzinfo=timezone.utc).astimezone(tz=None)
    time = time.replace(tzinfo=None)
34 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 49 50 51
        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]
52 53 54 55 56 57 58
        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
59
    text_buffer.add_message(
60 61 62 63 64 65 66 67 68 69 70
        txt=text,
        time=time,
        nickname=nick,
        nick_color=color,
        history=True,
        user=None,
        highlight=False,
        top=top,
        identifier=None,
        str_time=None,
        jid=None,
71 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 120 121 122 123
            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
124 125
    msg_count = 0
    msgs = []
126
    async for rsm in results:
127 128
        if top:
            for msg in rsm['mam']['results']:
Madhur Garg's avatar
Madhur Garg committed
129 130
                if msg['mam_result']['forwarded']['stanza'] \
                .xml.find('{%s}%s' % ('jabber:client', 'body')) is not None:
131
                    msgs.append(msg)
132
                if msg_count == amount:
133
                    tab.core.refresh_window()
Madhur Garg's avatar
Madhur Garg committed
134
                    return False
135 136 137 138 139 140
                msg_count += 1
            msgs.reverse()
            for msg in msgs:
                forwarded = msg['mam_result']['forwarded']
                timestamp = forwarded['delay']['stamp']
                message = forwarded['stanza']
141
                tab.last_stanza_id = msg['mam_result']['id']
142
                nick = str(message['from'])
143
                add_line(tab, text_buffer, message['body'], timestamp, nick, top)
144 145 146 147 148
        else:
            for msg in rsm['mam']['results']:
                forwarded = msg['mam_result']['forwarded']
                timestamp = forwarded['delay']['stamp']
                message = forwarded['stanza']
149
                nick = str(message['from'])
150
                add_line(tab, text_buffer, message['body'], timestamp, nick, top)
151
                tab.core.refresh_window()
Madhur Garg's avatar
Madhur Garg committed
152
    return False
153

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

164 165
    if amount >= 100:
        amount = 99
Madhur Garg's avatar
Madhur Garg committed
166 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

    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