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