1# Copyright (C) 2019 Philipp Hörist <philipp AT hoerist.com>
2#
3# This file is part of nbxmpp.
4#
5# This program is free software; you can redistribute it and/or
6# modify it under the terms of the GNU General Public License
7# as published by the Free Software Foundation; either version 3
8# of the License, or (at your option) any later version.
9#
10# This program is distributed in the hope that it will be useful,
11# but WITHOUT ANY WARRANTY; without even the implied warranty of
12# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13# GNU General Public License for more details.
14#
15# You should have received a copy of the GNU General Public License
16# along with this program; If not, see <http://www.gnu.org/licenses/>.
17
18from nbxmpp.protocol import Error as ErrorStanza
19from nbxmpp.protocol import ERR_BAD_REQUEST
20from nbxmpp.protocol import ERR_FEATURE_NOT_IMPLEMENTED
21from nbxmpp.protocol import NodeProcessed
22from nbxmpp.protocol import Iq
23from nbxmpp.namespaces import Namespace
24from nbxmpp.structs import StanzaHandler
25from nbxmpp.structs import IBBData
26from nbxmpp.util import b64decode
27from nbxmpp.util import b64encode
28from nbxmpp.modules.base import BaseModule
29from nbxmpp.modules.util import process_response
30from nbxmpp.task import iq_request_task
31
32
33class IBB(BaseModule):
34    def __init__(self, client):
35        BaseModule.__init__(self, client)
36
37        self._client = client
38        self.handlers = [
39            StanzaHandler(name='iq',
40                          callback=self._process_ibb,
41                          ns=Namespace.IBB,
42                          priority=20),
43        ]
44
45    def _process_ibb(self, _client, stanza, properties):
46        if properties.type.is_set:
47            open_ = stanza.getTag('open', namespace=Namespace.IBB)
48            if open_ is not None:
49                properties.ibb = self._parse_open(stanza, open_)
50                return
51
52            close = stanza.getTag('close', namespace=Namespace.IBB)
53            if close is not None:
54                properties.ibb = self._parse_close(stanza, close)
55                return
56
57            data = stanza.getTag('data', namespace=Namespace.IBB)
58            if data is not None:
59                properties.ibb = self._parse_data(stanza, data)
60                return
61
62    def _parse_open(self, stanza, open_):
63        attrs = open_.getAttrs()
64        try:
65            block_size = int(attrs.get('block-size'))
66        except Exception as error:
67            self._log.warning(error)
68            self._log.warning(stanza)
69            self._client.send_stanza(ErrorStanza(stanza, ERR_BAD_REQUEST))
70            raise NodeProcessed
71
72        if block_size > 65535:
73            self._log.warning('Invalid block-size')
74            self._client.send_stanza(ErrorStanza(stanza, ERR_BAD_REQUEST))
75            raise NodeProcessed
76
77        sid = attrs.get('sid')
78        if not sid:
79            self._log.warning('Invalid sid')
80            self._client.send_stanza(ErrorStanza(stanza, ERR_BAD_REQUEST))
81            raise NodeProcessed
82
83        type_ = attrs.get('stanza')
84        if type_ == 'message':
85            self._client.send_stanza(ErrorStanza(stanza,
86                                                 ERR_FEATURE_NOT_IMPLEMENTED))
87            raise NodeProcessed
88
89        return IBBData(type='open', block_size=block_size, sid=sid)
90
91    def _parse_close(self, stanza, close):
92        sid = close.getAttrs().get('sid')
93        if sid is None:
94            self._log.warning('Invalid sid')
95            self._client.send_stanza(ErrorStanza(stanza, ERR_BAD_REQUEST))
96            raise NodeProcessed
97        return IBBData(type='close', sid=sid)
98
99    def _parse_data(self, stanza, data):
100        attrs = data.getAttrs()
101
102        sid = attrs.get('sid')
103        if sid is None:
104            self._log.warning('Invalid sid')
105            self._client.send_stanza(ErrorStanza(stanza, ERR_BAD_REQUEST))
106            raise NodeProcessed
107
108        try:
109            seq = int(attrs.get('seq'))
110        except Exception:
111            self._log.exception('Invalid seq')
112            self._client.send_stanza(ErrorStanza(stanza, ERR_BAD_REQUEST))
113            raise NodeProcessed
114
115        try:
116            decoded_data = b64decode(data.getData(), return_type=bytes)
117        except Exception:
118            self._log.exception('Failed to decode IBB data')
119            self._client.send_stanza(ErrorStanza(stanza, ERR_BAD_REQUEST))
120            raise NodeProcessed
121
122        return IBBData(type='data', sid=sid, seq=seq, data=decoded_data)
123
124    def send_reply(self, stanza, error=None):
125        if error is None:
126            reply = stanza.buildReply('result')
127            reply.getChildren().clear()
128        else:
129            reply = ErrorStanza(stanza, error)
130        self._client.send_stanza(reply)
131
132    @iq_request_task
133    def send_open(self, jid, sid, block_size):
134        _task = yield
135
136        response = yield _make_ibb_open(jid, sid, block_size)
137        yield process_response(response)
138
139    @iq_request_task
140    def send_close(self, jid, sid):
141        _task = yield
142
143        response = yield _make_ibb_close(jid, sid)
144        yield process_response(response)
145
146    @iq_request_task
147    def send_data(self, jid, sid, seq, data):
148        _task = yield
149
150        response = yield _make_ibb_data(jid, sid, seq, data)
151        yield process_response(response)
152
153
154def _make_ibb_open(jid, sid, block_size):
155    iq = Iq('set', to=jid)
156    iq.addChild('open',
157                {'block-size': block_size, 'sid': sid, 'stanza': 'iq'},
158                namespace=Namespace.IBB)
159    return iq
160
161
162def _make_ibb_close(jid, sid):
163    iq = Iq('set', to=jid)
164    iq.addChild('close', {'sid': sid}, namespace=Namespace.IBB)
165    return iq
166
167
168def _make_ibb_data(jid, sid, seq, data):
169    iq = Iq('set', to=jid)
170    ibb_data = iq.addChild('data',
171                           {'sid': sid, 'seq': seq},
172                           namespace=Namespace.IBB)
173    ibb_data.setData(b64encode(data))
174    return iq
175