Add support for XEP-0313: Message Archive Management

NOTE: XEP-0313 is still very experimental, and there will likely be
      API changes in the future.

5 files changed, 242 insertions, 0 deletions

diff --git a/setup.py b/setup.py
index 3ed1a3a9..55c73f59 100755
--- a/setup.py
+++ b/setup.py
@@ -105,6 +105,7 @@ packages     = [ 'sleekxmpp',
                  'sleekxmpp/plugins/xep_0280',
                  'sleekxmpp/plugins/xep_0297',
                  'sleekxmpp/plugins/xep_0308',
+                 'sleekxmpp/plugins/xep_0313',
                  'sleekxmpp/features',
                  'sleekxmpp/features/feature_mechanisms',
                  'sleekxmpp/features/feature_mechanisms/stanza',
diff --git a/sleekxmpp/plugins/__init__.py b/sleekxmpp/plugins/__init__.py
index 66027fbf..6188a65e 100644
--- a/sleekxmpp/plugins/__init__.py
+++ b/sleekxmpp/plugins/__init__.py
@@ -72,4 +72,5 @@ __all__ = [
     'xep_0297',  # Stanza Forwarding
     'xep_0302',  # XMPP Compliance Suites 2012
     'xep_0308',  # Last Message Correction
+    'xep_0313',  # Message Archive Management
 ]
diff --git a/sleekxmpp/plugins/xep_0313/__init__.py b/sleekxmpp/plugins/xep_0313/__init__.py
new file mode 100644
index 00000000..8b6ed97d
--- /dev/null
+++ b/sleekxmpp/plugins/xep_0313/__init__.py
@@ -0,0 +1,15 @@
+"""
+    SleekXMPP: The Sleek XMPP Library
+    Copyright (C) 2012 Nathanael C. Fritz, Lance J.T. Stout
+    This file is part of SleekXMPP.
+
+    See the file LICENSE for copying permissio
+"""
+
+from sleekxmpp.plugins.base import register_plugin
+
+from sleekxmpp.plugins.xep_0313.stanza import Result, MAM, Preferences
+from sleekxmpp.plugins.xep_0313.mam import XEP_0313
+
+
+register_plugin(XEP_0313)
diff --git a/sleekxmpp/plugins/xep_0313/mam.py b/sleekxmpp/plugins/xep_0313/mam.py
new file mode 100644
index 00000000..15aee828
--- /dev/null
+++ b/sleekxmpp/plugins/xep_0313/mam.py
@@ -0,0 +1,92 @@
+"""
+    SleekXMPP: The Sleek XMPP Library
+    Copyright (C) 2012 Nathanael C. Fritz, Lance J.T. Stout
+    This file is part of SleekXMPP.
+
+    See the file LICENSE for copying permissio
+"""
+
+import logging
+
+import sleekxmpp
+from sleekxmpp.stanza import Message, Iq
+from sleekxmpp.exceptions import XMPPError
+from sleekxmpp.xmlstream.handler import Collector
+from sleekxmpp.xmlstream.matcher import StanzaPath
+from sleekxmpp.xmlstream import register_stanza_plugin
+from sleekxmpp.plugins import BasePlugin
+from sleekxmpp.plugins.xep_0313 import stanza
+
+
+log = logging.getLogger(__name__)
+
+
+class XEP_0313(BasePlugin):
+
+    """
+    XEP-0313 Message Archive Management
+    """
+
+    name = 'xep_0313'
+    description = 'XEP-0313: Message Archive Management'
+    dependencies = set(['xep_0030', 'xep_0050', 'xep_0059', 'xep_0297'])
+    stanza = stanza
+
+    def plugin_init(self):
+        register_stanza_plugin(Iq, stanza.MAM)
+        register_stanza_plugin(Iq, stanza.Preferences)
+        register_stanza_plugin(Message, stanza.Result)
+        register_stanza_plugin(stanza.MAM, self.xmpp['xep_0059'].stanza.Set)
+
+    def retrieve(self, jid=None, start=None, end=None, with_jid=None, ifrom=None,
+                 block=True, timeout=None, callback=None, iterator=False):
+        iq = self.xmpp.Iq()
+        query_id = iq['id']
+
+        iq['to'] = jid
+        iq['from'] = ifrom
+        iq['type'] = 'get'
+        iq['mam']['queryid'] = query_id
+        iq['mam']['start'] = start
+        iq['mam']['end'] = end
+        iq['mam']['with'] = with_jid
+
+        collector = Collector(
+            'MAM_Results_%s' % query_id,
+            StanzaPath('message/mam_result@queryid=%s' % query_id))
+        self.xmpp.register_handler(collector)
+
+        if iterator:
+            return self.xmpp['xep_0059'].iterate(iq, 'mam', 'results')
+        elif not block and callback is not None:
+            def wrapped_cb(iq):
+                results = collector.stop()
+                if iq['type'] == 'result':
+                    iq['mam']['results'] = results
+                callback(iq)
+            return iq.send(block=block, timeout=timeout, callback=wrapped_cb)
+        else:
+            try:
+                resp = iq.send(block=block, timeout=timeout, callback=callback)
+                resp['mam']['results'] = collector.stop()
+                return resp
+            except XMPPError as e:
+                collector.stop()
+                raise e
+
+    def set_preferences(self, jid=None, default=None, always=None, never=None,
+                        ifrom=None, block=True, timeout=None, callback=None):
+        iq = self.xmpp.Iq()
+        iq['type'] = 'set'
+        iq['to'] = jid
+        iq['from'] = ifrom
+        iq['mam_prefs']['default'] = default
+        iq['mam_prefs']['always'] = always
+        iq['mam_prefs']['never'] = never
+        return iq.send(block=block, timeout=timeout, callback=callback)
+
+    def get_configuration_commands(self, jid, **kwargs):
+        return self.xmpp['xep_0030'].get_items(
+                jid=jid,
+                node='urn:xmpp:mam#configure',
+                **kwargs)
diff --git a/sleekxmpp/plugins/xep_0313/stanza.py b/sleekxmpp/plugins/xep_0313/stanza.py
new file mode 100644
index 00000000..6552dc5a
--- /dev/null
+++ b/sleekxmpp/plugins/xep_0313/stanza.py
@@ -0,0 +1,133 @@
+"""
+    SleekXMPP: The Sleek XMPP Library
+    Copyright (C) 2012 Nathanael C. Fritz, Lance J.T. Stout
+    This file is part of SleekXMPP.
+
+    See the file LICENSE for copying permissio
+"""
+
+import datetime as dt
+
+from sleekxmpp.jid import JID
+from sleekxmpp.xmlstream import ElementBase, ET
+from sleekxmpp.plugins import xep_0082
+
+
+class MAM(ElementBase):
+    name = 'query'
+    namespace = 'urn:xmpp:mam:tmp'
+    plugin_attrib = 'mam'
+    interfaces = set(['queryid', 'start', 'end', 'with', 'results'])
+    sub_interfaces = set(['start', 'end', 'with'])
+
+    def setup(self, xml=None):
+        ElementBase.setup(self, xml)
+
+        self._results = []
+
+    def get_start(self):
+        timestamp = self._get_attr('start')
+        return xep_0082.parse(timestamp)
+
+    def set_start(self, value):
+        if isinstance(value, dt.datetime):
+            value = xep_0082.format_datetime(value)
+        self._set_attr('start', value)
+
+    def get_end(self):
+        timestamp = self._get_sub_text('end')
+        return xep_0082.parse(timestamp)
+
+    def set_end(self, value):
+        if isinstance(value, dt.datetime):
+            value = xep_0082.format_datetime(value)
+        self._set_sub_text('end', value)
+
+    def get_with(self):
+        return JID(self._get_sub_text('with'))
+
+    def set_with(self, value):
+        self._set_sub_text('with', str(value))
+
+
+    # The results interface is meant only as an easy
+    # way to access the set of collected message responses
+    # from the query.
+
+    def get_results(self):
+        return self._results
+
+    def set_results(self, values):
+        self._results = values
+
+    def del_results(self):
+        self._results = []
+
+
+class Preferences(ElementBase):
+    name = 'prefs'
+    namespace = 'urn:xmpp:mam:tmp'
+    plugin_attrib = 'mam_prefs'
+    interfaces = set(['default', 'always', 'never'])
+    sub_interfaces = set(['always', 'never'])
+
+    def get_always(self):
+        results = set()
+
+        jids = self.xml.findall('{%s}always/{%s}jid' % (
+            self.namespace, self.namespace))
+
+        for jid in jids:
+            results.add(JID(jid.text))
+
+        return results
+
+    def set_always(self, value):
+        self._set_sub_text('always', '', keep=True)
+        always = self.xml.find('{%s}always' % self.namespace)
+        always.clear()
+
+        if not isinstance(value, (list, set)):
+            value = [value]
+
+        for jid in value:
+            jid_xml = ET.Element('{%s}jid' % self.namespace)
+            jid_xml.text = str(jid)
+            always.append(jid_xml)
+
+    def get_never(self):
+        results = set()
+
+        jids = self.xml.findall('{%s}never/{%s}jid' % (
+            self.namespace, self.namespace))
+
+        for jid in jids:
+            results.add(JID(jid.text))
+
+        return results
+
+    def set_never(self, value):
+        self._set_sub_text('never', '', keep=True)
+        never = self.xml.find('{%s}never' % self.namespace)
+        never.clear()
+
+        if not isinstance(value, (list, set)):
+            value = [value]
+
+        for jid in value:
+            jid_xml = ET.Element('{%s}jid' % self.namespace)
+            jid_xml.text = str(jid)
+            never.append(jid_xml)
+
+
+class Result(ElementBase):
+    name = 'result'
+    namespace = 'urn:xmpp:mam:tmp'
+    plugin_attrib = 'mam_result'
+    interfaces = set(['forwarded', 'queryid', 'id'])
+
+    def get_forwarded(self):
+        return self.parent()['forwarded']
+
+    def del_forwarded(self):
+        del self.parent()['forwarded']