diff --git a/examples/http_upload.py b/examples/http_upload.py index cbe4ab3f..e34b7b01 100755 --- a/examples/http_upload.py +++ b/examples/http_upload.py @@ -9,20 +9,13 @@ See the file LICENSE for copying permission. """ -import sys - import logging from getpass import getpass from argparse import ArgumentParser import slixmpp -from slixmpp.exceptions import XMPPError, IqError from slixmpp import asyncio -from urllib.parse import urlparse -from http.client import HTTPConnection, HTTPSConnection -from mimetypes import MimeTypes - log = logging.getLogger(__name__) @@ -36,65 +29,19 @@ class HttpUpload(slixmpp.ClientXMPP): slixmpp.ClientXMPP.__init__(self, jid, password) self.recipient = recipient - self.file = open(filename, 'rb') - self.size = self.file.seek(0, 2) - self.file.seek(0) - self.content_type = MimeTypes().guess_type(filename)[0] or 'application/octet-stream' + self.filename = filename self.add_event_handler("session_start", self.start) @asyncio.coroutine def start(self, event): - log.info('Uploading file %s...', self.file.name) - - info_iq = yield from self['xep_0363'].find_upload_service() - 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 = 'Uploaded Image' % get - else: - html = '%s' % (get, get) + log.info('Uploading file %s...', self.filename) + url = yield from self['xep_0363'].upload_file(self.filename) + log.info('Upload success!') log.info('Sending file to %s', self.recipient) - self.send_message(self.recipient, get, mhtml=html) + html = '%s' % (url, url) + self.send_message(self.recipient, url, mhtml=html) self.disconnect() diff --git a/slixmpp/plugins/xep_0363/http_upload.py b/slixmpp/plugins/xep_0363/http_upload.py index 3026a8c9..9218600b 100644 --- a/slixmpp/plugins/xep_0363/http_upload.py +++ b/slixmpp/plugins/xep_0363/http_upload.py @@ -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'] diff --git a/slixmpp/plugins/xep_0363/stanza.py b/slixmpp/plugins/xep_0363/stanza.py index 4795f96d..a57c5fb0 100644 --- a/slixmpp/plugins/xep_0363/stanza.py +++ b/slixmpp/plugins/xep_0363/stanza.py @@ -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 = ''