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