mam.py 8.82 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).
"""

mathieui's avatar
mathieui committed
9 10
import asyncio
import logging
11
import random
12
from datetime import datetime, timedelta, timezone
13
from hashlib import md5
mathieui's avatar
mathieui committed
14 15 16 17 18 19 20
from typing import (
    AsyncIterable,
    Callable,
    Dict,
    List,
    Optional,
)
21

mathieui's avatar
mathieui committed
22
from slixmpp import JID, Message as SMessage
23
from slixmpp.exceptions import IqError, IqTimeout
24
from poezio.theming import get_theme
25
from poezio import tabs
26
from poezio import xhtml, colors
27
from poezio.config import config
mathieui's avatar
mathieui committed
28 29
from poezio.text_buffer import TextBuffer, HistoryGap
from poezio.ui.types import BaseMessage, Message
Madhur Garg's avatar
Madhur Garg committed
30 31


mathieui's avatar
mathieui committed
32 33
log = logging.getLogger(__name__)

Madhur Garg's avatar
Madhur Garg committed
34 35 36 37
class DiscoInfoException(Exception): pass
class MAMQueryException(Exception): pass
class NoMAMSupportException(Exception): pass

38

mathieui's avatar
mathieui committed
39 40
def make_line(
        tab: tabs.Tab,
41 42 43
        text: str,
        time: datetime,
        nick: str,
mathieui's avatar
mathieui committed
44 45
        identifier: str = '',
    ) -> Message:
46 47
    """Adds a textual entry in the TextBuffer"""

48
    # Convert to local timezone
49 50
    time = time.replace(tzinfo=timezone.utc).astimezone(tz=None)
    time = time.replace(tzinfo=None)
51

Maxime Buquet's avatar
Maxime Buquet committed
52
    deterministic = config.get_by_tabname('deterministic_nick_colors', tab.jid.bare)
53 54 55
    if isinstance(tab, tabs.MucTab):
        nick = nick.split('/')[1]
        user = tab.get_user_by_name(nick)
56 57 58 59 60 61 62 63 64 65
        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
66
                    nick_pos = int(md5(nick.encode('utf-8')).hexdigest(), 16) % mod
67
                    color = theme.LIST_COLOR_NICKNAMES[nick_pos]
68 69 70 71 72
        else:
            color = random.choice(list(xhtml.colors))
            color = xhtml.colors.get(color)
            color = (color, -1)
    else:
mathieui's avatar
mathieui committed
73 74 75 76 77 78 79 80 81 82 83 84 85
        if nick.split('/')[0] == tab.core.xmpp.boundjid.bare:
            color = get_theme().COLOR_OWN_NICK
        else:
            color = get_theme().COLOR_REMOTE_USER
        nick = tab.get_nick()
    return Message(
        txt=text,
        identifier=identifier,
        time=time,
        nickname=nick,
        nick_color=color,
        history=True,
        user=None,
86 87
    )

Maxime Buquet's avatar
Maxime Buquet committed
88

mathieui's avatar
mathieui committed
89
async def get_mam_iterator(
Madhur Garg's avatar
Madhur Garg committed
90 91 92 93
        core,
        groupchat: bool,
        remote_jid: JID,
        amount: int,
mathieui's avatar
mathieui committed
94
        reverse: bool = True,
Madhur Garg's avatar
Madhur Garg committed
95 96 97
        start: Optional[datetime] = None,
        end: Optional[datetime] = None,
        before: Optional[str] = None,
mathieui's avatar
mathieui committed
98 99
    ) -> AsyncIterable[Message]:
    """Get an async iterator for this mam query"""
100
    try:
101 102
        query_jid = remote_jid if groupchat else JID(core.xmpp.boundjid.bare)
        iq = await core.xmpp.plugin['xep_0030'].get_info(jid=query_jid)
103
    except (IqError, IqTimeout):
mathieui's avatar
mathieui committed
104
        raise DiscoInfoException()
Madhur Garg's avatar
Madhur Garg committed
105
    if 'urn:xmpp:mam:2' not in iq['disco_info'].get_features():
mathieui's avatar
mathieui committed
106
        raise NoMAMSupportException()
Madhur Garg's avatar
Madhur Garg committed
107 108 109 110 111 112 113 114

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

    if groupchat:
        args['jid'] = remote_jid
115
    else:
Madhur Garg's avatar
Madhur Garg committed
116 117
        args['with_jid'] = remote_jid

mathieui's avatar
mathieui committed
118 119 120 121 122
    if amount > 0:
        args['rsm'] = {'max': amount}
    args['start'] = start
    args['end'] = end
    return core.xmpp['xep_0313'].retrieve(**args)
Madhur Garg's avatar
Madhur Garg committed
123 124


mathieui's avatar
mathieui committed
125 126 127 128 129 130 131 132 133 134
def _parse_message(msg: SMessage) -> Dict:
    """Parse info inside a MAM forwarded message"""
    forwarded = msg['mam_result']['forwarded']
    message = forwarded['stanza']
    return {
        'time': forwarded['delay']['stamp'],
        'nick': str(message['from']),
        'text': message['body'],
        'identifier': message['origin-id']
    }
Maxime Buquet's avatar
Maxime Buquet committed
135

Madhur Garg's avatar
Madhur Garg committed
136

mathieui's avatar
mathieui committed
137 138 139 140
async def retrieve_messages(tab: tabs.Tab,
                            results: AsyncIterable[SMessage],
                            amount: int = 100) -> List[Message]:
    """Run the MAM query and put messages in order"""
Madhur Garg's avatar
Madhur Garg committed
141
    text_buffer = tab._text_buffer
142 143
    msg_count = 0
    msgs = []
mathieui's avatar
mathieui committed
144 145 146 147
    to_add = []
    last_stanza_id = tab.last_stanza_id
    try:
        async for rsm in results:
148
            for msg in rsm['mam']['results']:
Madhur Garg's avatar
Madhur Garg committed
149
                if msg['mam_result']['forwarded']['stanza'] \
mathieui's avatar
mathieui committed
150 151 152 153 154
                        .xml.find('{%s}%s' % ('jabber:client', 'body')) is not None:
                    args = _parse_message(msg)
                    msgs.append(make_line(tab, **args))
            for msg in reversed(msgs):
                to_add.append(msg)
155
                msg_count += 1
mathieui's avatar
mathieui committed
156 157 158 159 160 161 162 163 164
                if msg_count == amount:
                    to_add.reverse()
                    return to_add
            msgs = []
        to_add.reverse()
        return to_add
    except (IqError, IqTimeout) as exc:
        log.debug('Unable to complete MAM query: %s', exc, exc_info=True)
        raise MAMQueryException('Query interrupted')
165

Maxime Buquet's avatar
Maxime Buquet committed
166

mathieui's avatar
mathieui committed
167 168 169 170
async def fetch_history(tab: tabs.Tab,
                        start: Optional[datetime] = None,
                        end: Optional[datetime] = None,
                        amount: Optional[int] = None) -> None:
171
    remote_jid = tab.jid
mathieui's avatar
mathieui committed
172 173 174 175 176 177
    if not end:
        for msg in tab._text_buffer.messages:
            if isinstance(msg, Message):
                end = msg.time
                end -= timedelta(microseconds=1)
                break
Madhur Garg's avatar
Madhur Garg committed
178 179
    if end is None:
        end = datetime.now()
180 181 182 183
    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
184

mathieui's avatar
mathieui committed
185 186 187 188
    if start is not None:
        start = start.replace(tzinfo=tzone).astimezone(tz=timezone.utc)
        start = start.replace(tzinfo=None)
        start = datetime.strftime(start, '%Y-%m-%dT%H:%M:%SZ')
Madhur Garg's avatar
Madhur Garg committed
189

mathieui's avatar
mathieui committed
190 191 192 193 194
    mam_iterator = await get_mam_iterator(
        core=tab.core,
        groupchat=isinstance(tab, tabs.MucTab),
        remote_jid=remote_jid,
        amount=amount,
Maxime Buquet's avatar
Maxime Buquet committed
195
        end=end,
mathieui's avatar
mathieui committed
196 197
        start=start,
        reverse=True,
Maxime Buquet's avatar
Maxime Buquet committed
198
    )
mathieui's avatar
mathieui committed
199
    return await retrieve_messages(tab, mam_iterator, amount)
Madhur Garg's avatar
Madhur Garg committed
200

mathieui's avatar
mathieui committed
201 202 203 204 205 206 207 208 209 210 211 212 213 214 215
async def fill_missing_history(tab: tabs.Tab, gap: HistoryGap) -> None:
    start = gap.last_timestamp_before_leave
    end = gap.first_timestamp_after_join
    if start:
        start = start + timedelta(seconds=1)
    if end:
        end = end - timedelta(seconds=1)
    try:
        messages = await fetch_history(tab, start=start, end=end, amount=999)
        tab._text_buffer.add_history_messages(messages, gap=gap)
        tab.core.refresh_window()
    except (NoMAMSupportException, MAMQueryException, DiscoInfoException):
        return
    finally:
        tab.query_status = False
Maxime Buquet's avatar
Maxime Buquet committed
216

mathieui's avatar
mathieui committed
217 218
async def on_new_tab_open(tab: tabs.Tab) -> None:
    """Called when opening a new tab"""
Madhur Garg's avatar
Madhur Garg committed
219 220 221
    amount = 2 * tab.text_win.height
    end = datetime.now()
    for message in tab._text_buffer.messages:
mathieui's avatar
mathieui committed
222 223 224 225
        if isinstance(message, Message) and message.time < end:
            end = message.time
            break
    end = end - timedelta(microseconds=1)
Madhur Garg's avatar
Madhur Garg committed
226
    try:
mathieui's avatar
mathieui committed
227 228
        messages = await fetch_history(tab, end=end, amount=amount)
        tab._text_buffer.add_history_messages(messages)
Madhur Garg's avatar
Madhur Garg committed
229 230
    except (NoMAMSupportException, MAMQueryException, DiscoInfoException):
        return None
mathieui's avatar
mathieui committed
231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252
    finally:
        tab.query_status = False


def schedule_tab_open(tab: tabs.Tab) -> None:
    """Set the query status and schedule a MAM query"""
    tab.query_status = True
    asyncio.ensure_future(on_tab_open(tab))


async def on_tab_open(tab: tabs.Tab) -> None:
    gap = tab._text_buffer.find_last_gap_muc()
    if gap is not None:
        await fill_missing_history(tab, gap)
    else:
        await on_new_tab_open(tab)


def schedule_scroll_up(tab: tabs.Tab) -> None:
    """Set query status and schedule a scroll up"""
    tab.query_status = True
    asyncio.ensure_future(on_scroll_up(tab))
Madhur Garg's avatar
Madhur Garg committed
253

Maxime Buquet's avatar
Maxime Buquet committed
254

Madhur Garg's avatar
Madhur Garg committed
255
async def on_scroll_up(tab) -> None:
256 257
    tw = tab.text_win

258 259 260
    # 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.
261 262 263 264
    total, pos, height = len(tw.built_lines), tw.pos, tw.height
    rest = (total - pos) // height

    if rest > 1:
mathieui's avatar
mathieui committed
265
        tab.query_status = False
266 267
        return None

Madhur Garg's avatar
Madhur Garg committed
268
    try:
269 270 271
        # 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`).
mathieui's avatar
mathieui committed
272 273
        messages = await fetch_history(tab, amount=height)
        tab._text_buffer.add_history_messages(messages)
Madhur Garg's avatar
Madhur Garg committed
274 275 276 277 278 279
    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
mathieui's avatar
mathieui committed
280 281
    finally:
        tab.query_status = False