146 lines
4.1 KiB
Python
146 lines
4.1 KiB
Python
"""
|
|
slixmpp: The Slick XMPP Library
|
|
Copyright (C) 2011 Nathanael C. Fritz
|
|
This file is part of slixmpp.
|
|
|
|
See the file LICENSE for copying permission.
|
|
"""
|
|
|
|
import logging
|
|
import zlib
|
|
|
|
|
|
from slixmpp.stanza import StreamFeatures
|
|
from slixmpp.xmlstream import RestartStream, register_stanza_plugin, ElementBase, StanzaBase
|
|
from slixmpp.xmlstream.matcher import *
|
|
from slixmpp.xmlstream.handler import *
|
|
from slixmpp.plugins import BasePlugin, register_plugin
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class Compression(ElementBase):
|
|
name = 'compression'
|
|
namespace = 'http://jabber.org/features/compress'
|
|
interfaces = {'methods'}
|
|
plugin_attrib = 'compression'
|
|
plugin_tag_map = {}
|
|
plugin_attrib_map = {}
|
|
|
|
def get_methods(self):
|
|
methods = []
|
|
for method in self.xml.findall('{%s}method' % self.namespace):
|
|
methods.append(method.text)
|
|
return methods
|
|
|
|
|
|
class Compress(StanzaBase):
|
|
name = 'compress'
|
|
namespace = 'http://jabber.org/protocol/compress'
|
|
interfaces = {'method'}
|
|
sub_interfaces = interfaces
|
|
plugin_attrib = 'compress'
|
|
plugin_tag_map = {}
|
|
plugin_attrib_map = {}
|
|
|
|
def setup(self, xml):
|
|
StanzaBase.setup(self, xml)
|
|
self.xml.tag = self.tag_name()
|
|
|
|
|
|
class Compressed(StanzaBase):
|
|
name = 'compressed'
|
|
namespace = 'http://jabber.org/protocol/compress'
|
|
interfaces = set()
|
|
plugin_tag_map = {}
|
|
plugin_attrib_map = {}
|
|
|
|
def setup(self, xml):
|
|
StanzaBase.setup(self, xml)
|
|
self.xml.tag = self.tag_name()
|
|
|
|
|
|
|
|
|
|
class ZlibSocket(object):
|
|
|
|
def __init__(self, socketobj):
|
|
self.__socket = socketobj
|
|
self.compressor = zlib.compressobj()
|
|
self.decompressor = zlib.decompressobj(zlib.MAX_WBITS)
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self.__socket, name)
|
|
|
|
def send(self, data):
|
|
sentlen = len(data)
|
|
data = self.compressor.compress(data)
|
|
data += self.compressor.flush(zlib.Z_SYNC_FLUSH)
|
|
log.debug(b'>>> (compressed)' + (data.encode("hex")))
|
|
#return self.__socket.send(data)
|
|
sentactuallen = self.__socket.send(data)
|
|
assert(sentactuallen == len(data))
|
|
|
|
return sentlen
|
|
|
|
def recv(self, *args, **kwargs):
|
|
data = self.__socket.recv(*args, **kwargs)
|
|
log.debug(b'<<< (compressed)' + data.encode("hex"))
|
|
return self.decompressor.decompress(self.decompressor.unconsumed_tail + data)
|
|
|
|
|
|
class XEP_0138(BasePlugin):
|
|
"""
|
|
XEP-0138: Compression
|
|
"""
|
|
name = "xep_0138"
|
|
description = "XEP-0138: Compression"
|
|
dependencies = {"xep_0030"}
|
|
|
|
def plugin_init(self):
|
|
self.xep = '0138'
|
|
self.description = 'Stream Compression (Generic)'
|
|
|
|
self.compression_methods = {'zlib': True}
|
|
|
|
register_stanza_plugin(StreamFeatures, Compression)
|
|
self.xmpp.register_stanza(Compress)
|
|
self.xmpp.register_stanza(Compressed)
|
|
|
|
self.xmpp.register_handler(
|
|
Callback('Compressed',
|
|
StanzaPath('compressed'),
|
|
self._handle_compressed,
|
|
instream=True))
|
|
|
|
self.xmpp.register_feature('compression',
|
|
self._handle_compression,
|
|
restart=True,
|
|
order=self.config.get('order', 5))
|
|
|
|
def register_compression_method(self, name, handler):
|
|
self.compression_methods[name] = handler
|
|
|
|
def _handle_compression(self, features):
|
|
for method in features['compression']['methods']:
|
|
if method in self.compression_methods:
|
|
log.info('Attempting to use %s compression' % method)
|
|
c = Compress(self.xmpp)
|
|
c['method'] = method
|
|
c.send(now=True)
|
|
return True
|
|
return False
|
|
|
|
def _handle_compressed(self, stanza):
|
|
self.xmpp.features.add('compression')
|
|
log.debug('Stream Compressed!')
|
|
compressed_socket = ZlibSocket(self.xmpp.socket)
|
|
self.xmpp.set_socket(compressed_socket)
|
|
raise RestartStream()
|
|
|
|
def _handle_failure(self, stanza):
|
|
pass
|
|
|
|
xep_0138 = XEP_0138
|
|
register_plugin(XEP_0138)
|