diff options
Diffstat (limited to 'slixmpp/plugins/xep_0363/http_upload.py')
| -rw-r--r-- | slixmpp/plugins/xep_0363/http_upload.py | 91 |
1 files changed, 82 insertions, 9 deletions
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'] |