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 typing import List
19from typing import Dict
20
21import hashlib
22from dataclasses import dataclass
23from dataclasses import asdict
24from dataclasses import field
25
26from nbxmpp.namespaces import Namespace
27from nbxmpp.protocol import NodeProcessed
28from nbxmpp.protocol import Node
29from nbxmpp.structs import StanzaHandler
30from nbxmpp.util import b64encode
31from nbxmpp.util import b64decode
32from nbxmpp.errors import MalformedStanzaError
33from nbxmpp.task import iq_request_task
34from nbxmpp.modules.base import BaseModule
35from nbxmpp.modules.util import raise_if_error
36from nbxmpp.modules.util import finalize
37
38
39class UserAvatar(BaseModule):
40
41    _depends = {
42        'publish': 'PubSub',
43        'request_item': 'PubSub',
44        'request_items': 'PubSub',
45        'delete': 'PubSub',
46    }
47
48    def __init__(self, client):
49        BaseModule.__init__(self, client)
50
51        self._client = client
52        self.handlers = [
53            StanzaHandler(name='message',
54                          callback=self._process_pubsub_avatar,
55                          ns=Namespace.PUBSUB_EVENT,
56                          priority=16),
57        ]
58
59    def _process_pubsub_avatar(self, _client, stanza, properties):
60        if not properties.is_pubsub_event:
61            return
62
63        if properties.pubsub_event.node != Namespace.AVATAR_METADATA:
64            return
65
66        item = properties.pubsub_event.item
67        if item is None:
68            # Retract, Deleted or Purged
69            return
70
71        metadata = item.getTag('metadata', namespace=Namespace.AVATAR_METADATA)
72        if metadata is None:
73            self._log.warning('No metadata node found')
74            self._log.warning(stanza)
75            raise NodeProcessed
76
77        if not metadata.getChildren():
78            self._log.info('Received avatar metadata: %s - no avatar set',
79                           properties.jid)
80            return
81
82        try:
83            data = AvatarMetaData.from_node(metadata, item.getAttr('id'))
84        except Exception as error:
85            self._log.warning('Malformed user avatar data: %s', error)
86            self._log.warning(stanza)
87            raise NodeProcessed
88
89        pubsub_event = properties.pubsub_event._replace(data=data)
90        self._log.info('Received avatar metadata: %s - %s',
91                       properties.jid, data)
92
93        properties.pubsub_event = pubsub_event
94
95    @iq_request_task
96    def request_avatar_data(self, id_, jid=None):
97        task = yield
98
99        item = yield self.request_item(Namespace.AVATAR_DATA,
100                                       id_=id_,
101                                       jid=jid)
102
103        raise_if_error(item)
104
105        if item is None:
106            yield task.set_result(None)
107
108        yield _get_avatar_data(item, id_)
109
110    @iq_request_task
111    def request_avatar_metadata(self, jid=None):
112        task = yield
113
114        items = yield self.request_items(Namespace.AVATAR_METADATA,
115                                         max_items=1,
116                                         jid=jid)
117
118        raise_if_error(items)
119
120        if not items:
121            yield task.set_result(None)
122
123        item = items[0]
124        metadata = item.getTag('metadata', namespace=Namespace.AVATAR_METADATA)
125        if metadata is None:
126            raise MalformedStanzaError('metadata node missing', item)
127
128        if not metadata.getChildren():
129            yield task.set_result(None)
130
131        yield AvatarMetaData.from_node(metadata, item.getAttr('id'))
132
133    @iq_request_task
134    def set_avatar(self, avatar, public=False):
135
136        task = yield
137
138        access_model = 'open' if public else 'presence'
139
140        if avatar is None:
141            result = yield self._publish_avatar_metadata(None, access_model)
142            raise_if_error(result)
143
144            result = yield self.delete(Namespace.AVATAR_DATA)
145            yield finalize(task, result)
146
147        result = yield self._publish_avatar(avatar, access_model)
148
149        yield finalize(task, result)
150
151    @iq_request_task
152    def _publish_avatar(self, avatar, access_model):
153        task = yield
154
155        options = {
156            'pubsub#persist_items': 'true',
157            'pubsub#access_model': access_model,
158        }
159
160        for info, data in avatar.pubsub_avatar_info():
161            item = _make_avatar_data_node(data)
162            self._log.info('Publish avatar data: %s, %s', info, access_model)
163
164            result = yield self.publish(Namespace.AVATAR_DATA,
165                                        item,
166                                        id_=info.id,
167                                        options=options,
168                                        force_node_options=True)
169
170            raise_if_error(result)
171
172        result = yield self._publish_avatar_metadata(avatar.metadata,
173                                                     access_model)
174
175        yield finalize(task, result)
176
177    @iq_request_task
178    def _publish_avatar_metadata(self, metadata, access_model):
179        task = yield
180
181        self._log.info('Publish avatar meta data: %s', metadata)
182
183        options = {
184            'pubsub#persist_items': 'true',
185            'pubsub#access_model': access_model,
186        }
187
188        if metadata is None:
189            metadata = AvatarMetaData()
190
191        result = yield self.publish(Namespace.AVATAR_METADATA,
192                                    metadata.to_node(),
193                                    id_=metadata.default,
194                                    options=options,
195                                    force_node_options=True)
196
197        yield finalize(task, result)
198
199    @iq_request_task
200    def set_access_model(self, public):
201        task = yield
202
203        access_model = 'open' if public else 'presence'
204
205        result = yield self._client.get_module('PubSub').set_access_model(
206            Namespace.AVATAR_DATA, access_model)
207
208        raise_if_error(result)
209
210        result = yield self._client.get_module('PubSub').set_access_model(
211            Namespace.AVATAR_METADATA, access_model)
212
213        yield finalize(task, result)
214
215
216def _get_avatar_data(item, id_):
217    data_node = item.getTag('data', namespace=Namespace.AVATAR_DATA)
218    if data_node is None:
219        raise MalformedStanzaError('data node missing', item)
220
221    data = data_node.getData()
222    if not data:
223        raise MalformedStanzaError('data node empty', item)
224
225    try:
226        avatar = b64decode(data, return_type=bytes)
227    except Exception as error:
228        raise MalformedStanzaError(f'decoding error: {error}', item)
229
230    avatar_sha = hashlib.sha1(avatar).hexdigest()
231    if avatar_sha != id_:
232        raise MalformedStanzaError(f'avatar does not match sha', item)
233
234    return AvatarData(data=avatar, sha=avatar_sha)
235
236
237def _make_metadata_node(infos):
238    item = Node('metadata', attrs={'xmlns': Namespace.AVATAR_METADATA})
239    for info in infos:
240        item.addChild('info', attrs=info.to_dict())
241    return item
242
243
244def _make_avatar_data_node(avatar):
245    item = Node('data', attrs={'xmlns': Namespace.AVATAR_DATA})
246    item.setData(b64encode(avatar.data))
247    return item
248
249
250def _get_info_attrs(avatar, avatar_sha, height, width):
251    info_attrs = {
252        'id': avatar_sha,
253        'bytes': len(avatar),
254        'type': 'image/png',
255    }
256
257    if height is not None:
258        info_attrs['height'] = height
259
260    if width is not None:
261        info_attrs['width'] = width
262
263    return info_attrs
264
265
266@dataclass
267class AvatarInfo:
268    bytes: int
269    id: str
270    type: str
271    url: str = None
272    height: int = None
273    width: int = None
274
275    def __post_init__(self):
276        if self.bytes is None:
277            raise ValueError
278        if self.id is None:
279            raise ValueError
280        if self.type is None:
281            raise ValueError
282
283        self.bytes = int(self.bytes)
284
285        if self.height is not None:
286            self.height = int(self.height)
287        if self.width is not None:
288            self.width = int(self.width)
289
290    def to_dict(self):
291        info_dict = asdict(self)
292        if self.height is None:
293            info_dict.pop('height')
294        if self.width is None:
295            info_dict.pop('width')
296        if self.url is None:
297            info_dict.pop('url')
298        return info_dict
299
300    def __hash__(self):
301        return hash(self.id)
302
303
304@dataclass
305class AvatarData:
306    data: bytes
307    sha: str
308
309
310@dataclass
311class AvatarMetaData:
312    infos: List[AvatarInfo] = field(default_factory=list)
313    default: AvatarInfo = None
314
315    @classmethod
316    def from_node(cls, node, default=None):
317        infos = []
318        info_nodes = node.getTags('info')
319        for info in info_nodes:
320            infos.append(AvatarInfo(
321                bytes=info.getAttr('bytes'),
322                id=info.getAttr('id'),
323                type=info.getAttr('type'),
324                url=info.getAttr('url'),
325                height=info.getAttr('height'),
326                width=info.getAttr('width')
327            ))
328        return cls(infos=infos, default=default)
329
330    def add_avatar_info(self, avatar_info, make_default=False):
331        self.infos.append(avatar_info)
332        if make_default:
333            self.default = avatar_info.id
334
335    def to_node(self):
336        return _make_metadata_node(self.infos)
337
338    @property
339    def avatar_shas(self):
340        return [info.id for info in self.infos]
341
342
343@dataclass
344class Avatar:
345    metadata: AvatarMetaData = field(default_factory=AvatarMetaData)
346    data: Dict[AvatarInfo, bytes] = field(init=False, default_factory=dict)
347
348    def add_image_source(self,
349                         data,
350                         type_,
351                         height,
352                         width,
353                         url=None,
354                         make_default=True):
355
356        sha = hashlib.sha1(data).hexdigest()
357        info = AvatarInfo(bytes=len(data),
358                          id=sha,
359                          type=type_,
360                          height=height,
361                          width=width,
362                          url=url)
363        self.metadata.add_avatar_info(info, make_default=make_default)
364        self.data[info] = AvatarData(data=data, sha=sha)
365
366    def pubsub_avatar_info(self):
367        for info, data in self.data.items():
368            if info.url is not None:
369                continue
370            yield info, data
371