1""" 2 slixmpp: The Slick XMPP Library 3 Copyright (C) 2018 Emmanuel Gil Peyrot 4 This file is part of slixmpp. 5 6 See the file LICENSE for copying permission. 7""" 8 9import logging 10import os.path 11 12from aiohttp import ClientSession 13from mimetypes import guess_type 14 15from slixmpp import Iq, __version__ 16from slixmpp.plugins import BasePlugin 17from slixmpp.xmlstream import register_stanza_plugin 18from slixmpp.xmlstream.handler import Callback 19from slixmpp.xmlstream.matcher import StanzaPath 20from slixmpp.plugins.xep_0363 import stanza, Request, Slot, Put, Get, Header 21 22log = logging.getLogger(__name__) 23 24class FileUploadError(Exception): 25 pass 26 27class UploadServiceNotFound(FileUploadError): 28 pass 29 30class FileTooBig(FileUploadError): 31 def __str__(self): 32 return 'File size too large: {} (max: {} bytes)' \ 33 .format(self.args[0], self.args[1]) 34 35class HTTPError(FileUploadError): 36 def __str__(self): 37 return 'Could not upload file: %d (%s)' % (self.args[0], self.args[1]) 38 39class XEP_0363(BasePlugin): 40 ''' This plugin only supports Python 3.5+ ''' 41 42 name = 'xep_0363' 43 description = 'XEP-0363: HTTP File Upload' 44 dependencies = {'xep_0030', 'xep_0128'} 45 stanza = stanza 46 default_config = { 47 'upload_service': None, 48 'max_file_size': float('+inf'), 49 'default_content_type': 'application/octet-stream', 50 } 51 52 def plugin_init(self): 53 register_stanza_plugin(Iq, Request) 54 register_stanza_plugin(Iq, Slot) 55 register_stanza_plugin(Slot, Put) 56 register_stanza_plugin(Slot, Get) 57 register_stanza_plugin(Put, Header, iterable=True) 58 59 self.xmpp.register_handler( 60 Callback('HTTP Upload Request', 61 StanzaPath('iq@type=get/http_upload_request'), 62 self._handle_request)) 63 64 def plugin_end(self): 65 self._http_session.close() 66 self.xmpp.remove_handler('HTTP Upload Request') 67 self.xmpp.remove_handler('HTTP Upload Slot') 68 self.xmpp['xep_0030'].del_feature(feature=Request.namespace) 69 70 def session_bind(self, jid): 71 self.xmpp.plugin['xep_0030'].add_feature(Request.namespace) 72 73 def _handle_request(self, iq): 74 self.xmpp.event('http_upload_request', iq) 75 76 async def find_upload_service(self, domain=None, timeout=None): 77 results = await self.xmpp['xep_0030'].get_info_from_domain( 78 domain=domain, timeout=timeout) 79 80 candidates = [] 81 for info in results: 82 for identity in info['disco_info']['identities']: 83 if identity[0] == 'store' and identity[1] == 'file': 84 candidates.append(info) 85 for info in candidates: 86 for feature in info['disco_info']['features']: 87 if feature == Request.namespace: 88 return info 89 90 def request_slot(self, jid, filename, size, content_type=None, ifrom=None, 91 timeout=None, callback=None, timeout_callback=None): 92 iq = self.xmpp.Iq() 93 iq['to'] = jid 94 iq['from'] = ifrom 95 iq['type'] = 'get' 96 request = iq['http_upload_request'] 97 request['filename'] = filename 98 request['size'] = str(size) 99 request['content-type'] = content_type or self.default_content_type 100 return iq.send(timeout=timeout, callback=callback, 101 timeout_callback=timeout_callback) 102 103 async def upload_file(self, filename, size=None, content_type=None, *, 104 input_file=None, ifrom=None, domain=None, timeout=None, 105 callback=None, timeout_callback=None): 106 ''' Helper function which does all of the uploading process. ''' 107 if self.upload_service is None: 108 info_iq = await self.find_upload_service( 109 domain=domain, timeout=timeout) 110 if info_iq is None: 111 raise UploadServiceNotFound() 112 self.upload_service = info_iq['from'] 113 for form in info_iq['disco_info'].iterables: 114 values = form['values'] 115 if values['FORM_TYPE'] == ['urn:xmpp:http:upload:0']: 116 try: 117 self.max_file_size = int(values['max-file-size']) 118 except (TypeError, ValueError): 119 log.error('Invalid max size received from HTTP File Upload service') 120 self.max_file_size = float('+inf') 121 break 122 123 if input_file is None: 124 input_file = open(filename, 'rb') 125 126 if size is None: 127 size = input_file.seek(0, 2) 128 input_file.seek(0) 129 130 if size > self.max_file_size: 131 raise FileTooBig(size, self.max_file_size) 132 133 if content_type is None: 134 content_type = guess_type(filename)[0] 135 if content_type is None: 136 content_type = self.default_content_type 137 138 basename = os.path.basename(filename) 139 slot_iq = await self.request_slot(self.upload_service, basename, size, 140 content_type, ifrom, timeout, 141 callback=callback, 142 timeout_callback=timeout_callback) 143 slot = slot_iq['http_upload_slot'] 144 145 headers = { 146 'Content-Length': str(size), 147 'Content-Type': content_type or self.default_content_type, 148 **{header['name']: header['value'] for header in slot['put']['headers']} 149 } 150 151 # Do the actual upload here. 152 async with ClientSession(headers={'User-Agent': 'slixmpp ' + __version__}) as session: 153 response = await session.put( 154 slot['put']['url'], 155 data=input_file, 156 headers=headers, 157 timeout=timeout) 158 if response.status >= 400: 159 raise HTTPError(response.status, await response.text()) 160 log.info('Response code: %d (%s)', response.status, await response.text()) 161 response.close() 162 return slot['get']['url'] 163