Simplify usage of HTTP File Upload plugin.

This makes it usable only on Python 3.5, as documented.
This commit is contained in:
Emmanuel Gil Peyrot
2018-03-08 14:28:55 +01:00
parent 29faf114a7
commit bd63b1ce70
3 changed files with 98 additions and 68 deletions

View File

@@ -9,7 +9,10 @@
import asyncio
import logging
from slixmpp import Iq
from aiohttp import ClientSession
from mimetypes import MimeTypes
from slixmpp import Iq, __version__
from slixmpp.plugins import BasePlugin
from slixmpp.xmlstream import register_stanza_plugin
from slixmpp.xmlstream.handler import Callback
@@ -18,19 +21,34 @@ from slixmpp.plugins.xep_0363 import stanza, Request, Slot, Put, Get, Header
log = logging.getLogger(__name__)
class FileUploadError(Exception):
pass
class UploadServiceNotFound(FileUploadError):
pass
class FileTooBig(FileUploadError):
pass
class XEP_0363(BasePlugin):
''' This plugin only supports Python 3.5+ '''
name = 'xep_0363'
description = 'XEP-0363: HTTP File Upload'
dependencies = {'xep_0030'}
dependencies = {'xep_0030', 'xep_0128'}
stanza = stanza
default_config = {
'upload_service': None,
'maximum_size': float('+inf'),
'default_content_type': 'application/octet-stream',
}
def plugin_init(self):
register_stanza_plugin(Iq, Request)
register_stanza_plugin(Iq, Slot)
register_stanza_plugin(Slot, Put)
register_stanza_plugin(Slot, Get)
register_stanza_plugin(Put, Header)
register_stanza_plugin(Put, Header, iterable=True)
self.xmpp.register_handler(
Callback('HTTP Upload Request',
@@ -38,6 +56,7 @@ class XEP_0363(BasePlugin):
self._handle_request))
def plugin_end(self):
self._http_session.close()
self.xmpp.remove_handler('HTTP Upload Request')
self.xmpp.remove_handler('HTTP Upload Slot')
self.xmpp['xep_0030'].del_feature(feature=Request.namespace)
@@ -48,15 +67,14 @@ class XEP_0363(BasePlugin):
def _handle_request(self, iq):
self.xmpp.event('http_upload_request', iq)
@asyncio.coroutine
def find_upload_service(self, ifrom=None, timeout=None, callback=None,
timeout_callback=None):
async def find_upload_service(self, ifrom=None, timeout=None, callback=None,
timeout_callback=None):
infos = [self.xmpp['xep_0030'].get_info(self.xmpp.boundjid.domain)]
iq_items = yield from self.xmpp['xep_0030'].get_items(
iq_items = await self.xmpp['xep_0030'].get_items(
self.xmpp.boundjid.domain, timeout=timeout)
items = iq_items['disco_items']['items']
infos += [self.xmpp['xep_0030'].get_info(item[0]) for item in items]
info_futures, _ = yield from asyncio.wait(infos, timeout=timeout)
info_futures, _ = await asyncio.wait(infos, timeout=timeout)
for future in info_futures:
info = future.result()
for identity in info['disco_info']['identities']:
@@ -72,6 +90,61 @@ class XEP_0363(BasePlugin):
request = iq['http_upload_request']
request['filename'] = filename
request['size'] = str(size)
request['content-type'] = content_type
request['content-type'] = content_type or self.default_content_type
return iq.send(timeout=timeout, callback=callback,
timeout_callback=timeout_callback)
async def upload_file(self, filename, size=None, content_type=None, *,
input_file=None, ifrom=None, timeout=None,
callback=None, timeout_callback=None):
''' Helper function which does all of the uploading process. '''
if self.upload_service is None:
info_iq = await self.find_upload_service(ifrom=ifrom, timeout=timeout)
if info_iq is None:
raise UploadServiceNotFound()
self.upload_service = info_iq['from']
for form in info_iq['disco_info'].iterables:
values = form['values']
if values['FORM_TYPE'] == ['urn:xmpp:http:upload:0']:
try:
self.max_file_size = int(values['max-file-size'])
except (TypeError, ValueError):
log.error('Invalid max size received from HTTP File Upload service')
self.max_file_size = float('+inf')
break
if input_file is None:
input_file = open(filename, 'rb')
if size is None:
size = input_file.seek(0, 2)
input_file.seek(0)
if size > self.max_file_size:
raise FileTooBig()
if content_type is None:
content_type = MimeTypes().guess_type(filename)[0]
if content_type is None:
content_type = self.default_content_type
slot_iq = await self.request_slot(self.upload_service, filename, size,
content_type, ifrom, timeout)
slot = slot_iq['http_upload_slot']
headers = {
'Content-Length': str(size),
'Content-Type': content_type or self.default_content_type,
**{header['name']: header['value'] for header in slot['put']['headers']}
}
# Do the actual upload here.
async with ClientSession(headers={'User-Agent': 'slixmpp ' + __version__}) as session:
response = await session.put(
slot['put']['url'],
data=input_file,
headers=headers,
timeout=timeout)
log.info('Response code: %d (%s)', response.status, await response.text())
response.close()
return slot['get']['url']

View File

@@ -35,4 +35,14 @@ class Header(ElementBase):
plugin_attrib = 'header'
name = 'header'
namespace = 'urn:xmpp:http:upload:0'
plugin_multi_attrib = 'headers'
interfaces = {'name', 'value'}
def get_value(self):
return self.xml.text
def set_value(self, value):
self.xml.text = value
def del_value(self):
self.xml.text = ''