mam.py 5.99 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 10
import asyncio
from datetime import datetime, timedelta, timezone
11
from slixmpp.exceptions import IqError, IqTimeout
12
from poezio.theming import get_theme
13
from poezio import tabs
14 15
from poezio.text_buffer import Message, TextBuffer

16
def add_line(text_buffer: TextBuffer, text: str, str_time: str, nick: str, top: bool):
17 18 19 20
    """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')
21 22
    time = time.replace(tzinfo=timezone.utc).astimezone(tz=None)
    time = time.replace(tzinfo=None)
23
    if nick is 'MAM':
24
        color = get_theme().COLOR_ME_MESSAGE
25 26
    else:
        color = get_theme().COLOR_OWN_NICK
27 28 29 30 31 32 33 34 35 36 37 38 39 40
    text_buffer.add_message(
        text,
        time,
        nick,
        color,
        True,  # History
        None,  # User
        False,  # Highlight
        top, #Top
        None,  # Identifier
        None,  # str_time
        None,  # Jid
    )

41
async def query(self, remote_jid, start, end, top):
42 43 44 45
    self.remote_jid = remote_jid
    self.start_date = start
    self.end_date = end
    text_buffer = self._text_buffer
46 47 48 49
    try:
        iq = await self.core.xmpp.plugin['xep_0030'].get_info(jid=remote_jid)
    except (IqError, IqTimeout):
        return self.information('Failed to retrieve messages', 'Error')
50 51
    if 'urn:xmpp:mam:2' not in iq['disco_info'].get_features():
        return self.core.information("This MUC doesn't support MAM.", "Error")
52
    if top:
53
        if isinstance(self, tabs.MucTab):
54 55 56 57
            try:
                results = self.core.xmpp['xep_0313'].retrieve(jid=self.remote_jid,
                iterator=True, reverse=top, end=self.end_date)
            except (IqError, IqTimeout):
58
                return self.core.information('Failed to retrieve messages', 'Error')
59
        else:
60 61 62 63
            try:
                results = self.core.xmpp['xep_0313'].retrieve(with_jid=self.remote_jid,
                iterator=True, reverse=top, end=self.end_date)
            except (IqError, IqTimeout):
64
                return self.core.information('Failed to retrieve messages', 'Error')
65
    else:
66
        if 'conference' in list(iq['disco_info']['identities'])[0]:
67 68 69 70
            try:
                results = self.core.xmpp['xep_0313'].retrieve(jid=self.remote_jid,
                iterator=True, reverse=top, start=self.start_date, end=self.end_date)
            except (IqError, IqTimeout):
71
                return self.core.information('Failed to retrieve messages', 'Error')
72
        else:
73 74 75 76
            try:
                results = self.core.xmpp['xep_0313'].retrieve(with_jid=self.remote_jid,
                iterator=True, reverse=top, start=self.start_date, end=self.end_date)
            except (IqError, IqTimeout):
77
                return self.core.information('Failed to retrieve messages', 'Error')
78 79
    msg_count = 0
    msgs = []
80
    async for rsm in results:
81 82
        if top:
            for msg in rsm['mam']['results']:
83 84
                if msg['mam_result']['forwarded']['stanza'].xml.find(
                    '{%s}%s' % ('jabber:client', 'body')) is not None:
85
                    msgs.append(msg)
86
                if msg_count == 10:
87
                    self.query_id = 0
Madhur Garg's avatar
Madhur Garg committed
88
                    self.core.refresh_window()
89 90 91 92
                    return
                msg_count += 1
            msgs.reverse()
            for msg in msgs:
93 94
                if msg is msgs[0]:
                    timestamp = msg['mam_result']['forwarded']['delay']['stamp']
95
                    add_line(text_buffer, 'Start of MAM query: ', timestamp, 'MAM', top)
96 97 98
                forwarded = msg['mam_result']['forwarded']
                timestamp = forwarded['delay']['stamp']
                message = forwarded['stanza']
99 100 101 102 103
                nick = str(message['from'])
                if isinstance(self, tabs.MucTab):
                    nick = nick.split('/')[1]
                else:
                    nick = nick.split('/')[0]
104
                add_line(text_buffer, message['body'], timestamp, nick, top)
105 106
                if msg is msgs[len(msgs)-1]:
                    timestamp = msg['mam_result']['forwarded']['delay']['stamp']
107
                    add_line(text_buffer, 'End of MAM query: ', timestamp, 'MAM', top)
108 109 110 111 112 113
                self.text_win.scroll_up(len(self.text_win.built_lines))
        else:
            for msg in rsm['mam']['results']:
                forwarded = msg['mam_result']['forwarded']
                timestamp = forwarded['delay']['stamp']
                message = forwarded['stanza']
114 115 116 117 118
                nick = str(message['from'])
                if 'conference' in list(iq['disco_info']['identities'])[0]:
                    nick = nick.split('/')[1]
                else:
                    nick = nick.split('/')[0]
119
                add_line(text_buffer, message['body'], timestamp, nick, top)
120
                self.core.refresh_window()
121 122
    if len(msgs) == 0:
        return self.core.information('No more messages left to retrieve', 'Info')
123
    self.query_id = 0
124 125 126 127 128 129 130 131 132 133

def mam_scroll(self):
    remote_jid = self.jid
    text_buffer = self._text_buffer
    end = datetime.now()
    for message in text_buffer.messages:
        time = message.time
        if time < end:
            end = time
    end = end + timedelta(seconds=-1)
Madhur Garg's avatar
Madhur Garg committed
134 135 136
    tzone = datetime.now().astimezone().tzinfo
    end = end.replace(tzinfo=tzone).astimezone(tz=timezone.utc)
    end = end.replace(tzinfo=None)
137
    end = datetime.strftime(end, '%Y-%m-%dT%H:%M:%SZ')
138
    start = False
139
    top = True
Madhur Garg's avatar
Madhur Garg committed
140 141 142 143
    pos = self.text_win.pos
    self.text_win.pos += self.text_win.height - 1
    if self.text_win.pos + self.text_win.height > len(self.text_win.built_lines):
        asyncio.ensure_future(query(self, remote_jid, start, end, top))
144
        self.query_id = 1
Madhur Garg's avatar
Madhur Garg committed
145 146 147 148
        self.text_win.pos = len(self.text_win.built_lines) - self.text_win.height
        if self.text_win.pos < 0:
            self.text_win.pos = 0
    return self.text_win.pos != pos