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,20 +9,13 @@
See the file LICENSE for copying permission. See the file LICENSE for copying permission.
""" """
import sys
import logging import logging
from getpass import getpass from getpass import getpass
from argparse import ArgumentParser from argparse import ArgumentParser
import slixmpp import slixmpp
from slixmpp.exceptions import XMPPError, IqError
from slixmpp import asyncio from slixmpp import asyncio
from urllib.parse import urlparse
from http.client import HTTPConnection, HTTPSConnection
from mimetypes import MimeTypes
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
@ -36,65 +29,19 @@ class HttpUpload(slixmpp.ClientXMPP):
slixmpp.ClientXMPP.__init__(self, jid, password) slixmpp.ClientXMPP.__init__(self, jid, password)
self.recipient = recipient self.recipient = recipient
self.file = open(filename, 'rb') self.filename = filename
self.size = self.file.seek(0, 2)
self.file.seek(0)
self.content_type = MimeTypes().guess_type(filename)[0] or 'application/octet-stream'
self.add_event_handler("session_start", self.start) self.add_event_handler("session_start", self.start)
@asyncio.coroutine @asyncio.coroutine
def start(self, event): def start(self, event):
log.info('Uploading file %s...', self.file.name) log.info('Uploading file %s...', self.filename)
url = yield from self['xep_0363'].upload_file(self.filename)
info_iq = yield from self['xep_0363'].find_upload_service() log.info('Upload success!')
if info_iq is None:
log.error('No upload service found on this server')
self.disconnect()
return
for form in info_iq['disco_info'].iterables:
values = form['values']
if values['FORM_TYPE'] == ['urn:xmpp:http:upload:0']:
max_file_size = int(values['max-file-size'])
if self.size > max_file_size:
log.error('File size bigger than max allowed')
self.disconnect()
return
break
else:
log.warn('Impossible to find max-file-size, assuming infinite storage space')
log.info('Using service %s', info_iq['from'])
slot_iq = yield from self['xep_0363'].request_slot(
info_iq['from'], self.file.name, self.size, self.content_type)
put = slot_iq['http_upload_slot']['put']['url']
get = slot_iq['http_upload_slot']['get']['url']
# Now we got the two URLs, we can start uploading the HTTP file.
put_scheme, put_host, put_path, _, _, _ = urlparse(put)
Connection = {'http': HTTPConnection, 'https': HTTPSConnection}[put_scheme]
conn = Connection(put_host)
conn.putrequest('PUT', put_path)
for header, value in slot_iq['http_upload_slot']['put']['headers']:
conn.putheader(header, value)
conn.putheader('Content-Length', self.size)
conn.putheader('Content-Type', self.content_type)
conn.endheaders(self.file.read())
response = conn.getresponse()
if response.status >= 400:
log.error('Failed to upload file: %d %s', response.status, response.reason)
self.disconnect()
return
log.info('Upload success! %d %s', response.status, response.reason)
if self.content_type.startswith('image/'):
html = '<body xmlns="http://www.w3.org/1999/xhtml"><img src="%s" alt="Uploaded Image"/></body>' % get
else:
html = '<body xmlns="http://www.w3.org/1999/xhtml"><a href="%s">%s</a></body>' % (get, get)
log.info('Sending file to %s', self.recipient) log.info('Sending file to %s', self.recipient)
self.send_message(self.recipient, get, mhtml=html) html = '<body xmlns="http://www.w3.org/1999/xhtml"><a href="%s">%s</a></body>' % (url, url)
self.send_message(self.recipient, url, mhtml=html)
self.disconnect() self.disconnect()

View File

@ -9,7 +9,10 @@
import asyncio import asyncio
import logging 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.plugins import BasePlugin
from slixmpp.xmlstream import register_stanza_plugin from slixmpp.xmlstream import register_stanza_plugin
from slixmpp.xmlstream.handler import Callback 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__) log = logging.getLogger(__name__)
class FileUploadError(Exception):
pass
class UploadServiceNotFound(FileUploadError):
pass
class FileTooBig(FileUploadError):
pass
class XEP_0363(BasePlugin): class XEP_0363(BasePlugin):
''' This plugin only supports Python 3.5+ '''
name = 'xep_0363' name = 'xep_0363'
description = 'XEP-0363: HTTP File Upload' description = 'XEP-0363: HTTP File Upload'
dependencies = {'xep_0030'} dependencies = {'xep_0030', 'xep_0128'}
stanza = stanza stanza = stanza
default_config = {
'upload_service': None,
'maximum_size': float('+inf'),
'default_content_type': 'application/octet-stream',
}
def plugin_init(self): def plugin_init(self):
register_stanza_plugin(Iq, Request) register_stanza_plugin(Iq, Request)
register_stanza_plugin(Iq, Slot) register_stanza_plugin(Iq, Slot)
register_stanza_plugin(Slot, Put) register_stanza_plugin(Slot, Put)
register_stanza_plugin(Slot, Get) register_stanza_plugin(Slot, Get)
register_stanza_plugin(Put, Header) register_stanza_plugin(Put, Header, iterable=True)
self.xmpp.register_handler( self.xmpp.register_handler(
Callback('HTTP Upload Request', Callback('HTTP Upload Request',
@ -38,6 +56,7 @@ class XEP_0363(BasePlugin):
self._handle_request)) self._handle_request))
def plugin_end(self): def plugin_end(self):
self._http_session.close()
self.xmpp.remove_handler('HTTP Upload Request') self.xmpp.remove_handler('HTTP Upload Request')
self.xmpp.remove_handler('HTTP Upload Slot') self.xmpp.remove_handler('HTTP Upload Slot')
self.xmpp['xep_0030'].del_feature(feature=Request.namespace) self.xmpp['xep_0030'].del_feature(feature=Request.namespace)
@ -48,15 +67,14 @@ class XEP_0363(BasePlugin):
def _handle_request(self, iq): def _handle_request(self, iq):
self.xmpp.event('http_upload_request', iq) self.xmpp.event('http_upload_request', iq)
@asyncio.coroutine async def find_upload_service(self, ifrom=None, timeout=None, callback=None,
def find_upload_service(self, ifrom=None, timeout=None, callback=None, timeout_callback=None):
timeout_callback=None):
infos = [self.xmpp['xep_0030'].get_info(self.xmpp.boundjid.domain)] 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) self.xmpp.boundjid.domain, timeout=timeout)
items = iq_items['disco_items']['items'] items = iq_items['disco_items']['items']
infos += [self.xmpp['xep_0030'].get_info(item[0]) for item in 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: for future in info_futures:
info = future.result() info = future.result()
for identity in info['disco_info']['identities']: for identity in info['disco_info']['identities']:
@ -72,6 +90,61 @@ class XEP_0363(BasePlugin):
request = iq['http_upload_request'] request = iq['http_upload_request']
request['filename'] = filename request['filename'] = filename
request['size'] = str(size) 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, return iq.send(timeout=timeout, callback=callback,
timeout_callback=timeout_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' plugin_attrib = 'header'
name = 'header' name = 'header'
namespace = 'urn:xmpp:http:upload:0' namespace = 'urn:xmpp:http:upload:0'
plugin_multi_attrib = 'headers'
interfaces = {'name', 'value'} 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 = ''