Coverage for slidge/core/mixins/attachment.py: 93%
248 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-10 04:45 +0000
1from __future__ import annotations
3import base64
4import functools
5import io
6import logging
7from collections.abc import Collection, Sequence
8from datetime import datetime
9from itertools import chain
10from mimetypes import guess_extension, guess_type
11from pathlib import Path
12from typing import Any, cast
13from urllib.parse import quote as urlquote
14from uuid import uuid4
15from xml.etree import ElementTree as ET
17import thumbhash
18from PIL import Image, ImageOps
19from slixmpp import JID, Message
20from slixmpp.plugins.xep_0264.stanza import Thumbnail
21from slixmpp.plugins.xep_0447.stanza import StatelessFileSharing
23from ...db.avatar import avatar_cache
24from ...db.models import Attachment
25from ...util.types import LegacyAttachment, MessageReference
26from ..attachment_upload import is_temp_path
27from .message_text import TextMessageMixin
30class AttachmentMixin(TextMessageMixin):
31 def __set_sims(
32 self,
33 msg: Message,
34 uploaded_url: str,
35 attachment: LegacyAttachment,
36 stored: Attachment,
37 thumbnail: Thumbnail | None,
38 ) -> None:
39 if stored.sims is not None:
40 ref = self.xmpp.plugin["xep_0372"].stanza.Reference(
41 xml=ET.fromstring(stored.sims)
42 )
43 msg.append(ref)
44 return
46 if attachment.data or attachment.stream or attachment.path:
47 ref = self.xmpp.plugin["xep_0385"].get_sims(
48 attachment.path, # type:ignore[arg-type]
49 [uploaded_url],
50 attachment.content_type,
51 attachment.caption,
52 file=attachment.stream,
53 data=attachment.data,
54 )
55 else:
56 sims = self.xmpp.plugin["xep_0385"].stanza.Sims()
57 ref = self.xmpp.plugin["xep_0372"].stanza.Reference()
58 ref["uri"] = uploaded_url
59 ref["type"] = "data"
60 sims["sources"].append(ref)
61 ref = self.xmpp.plugin["xep_0372"].stanza.Reference()
62 ref.append(sims)
63 ref["type"] = "data"
64 sims.enable("file")
65 if attachment.content_type:
66 sims["file"]["media-type"] = attachment.content_type
67 if attachment.caption:
68 sims["file"]["desc"] = attachment.caption
69 if attachment.size:
70 sims["file"]["size"] = attachment.size
71 if attachment.name:
72 ref["sims"]["file"]["name"] = attachment.name
73 if attachment.aio_stream is not None:
74 # revove the date in case we have downloaded the file to disk
75 # because no size was provided by the legacy module
76 ref["sims"]["file"]["date"] = None
78 if thumbnail is not None:
79 ref["sims"]["file"].append(thumbnail)
81 stored.sims = str(ref)
82 msg.append(ref)
84 def __set_sfs(
85 self,
86 msg: Message,
87 uploaded_url: str,
88 attachment: LegacyAttachment,
89 stored: Attachment,
90 thumbnail: Thumbnail | None = None,
91 ) -> None:
92 if stored.sfs is not None:
93 msg.append(StatelessFileSharing(xml=ET.fromstring(stored.sfs)))
94 return
96 if attachment.path or attachment.data or attachment.stream:
97 sfs = self.xmpp.plugin["xep_0447"].get_sfs( # type:ignore
98 attachment.path,
99 [uploaded_url],
100 attachment.content_type,
101 attachment.caption,
102 data=attachment.data,
103 file=attachment.stream,
104 )
105 else:
106 sfs = self.xmpp.plugin["xep_0447"].stanza.StatelessFileSharing()
107 ref = self.xmpp.plugin["xep_0447"].stanza.UrlData()
108 ref["target"] = uploaded_url
109 sfs["sources"].append(ref)
110 sfs.enable("file")
111 if attachment.content_type:
112 sfs["file"]["media-type"] = attachment.content_type
113 if attachment.name:
114 sfs["file"]["name"] = attachment.name
115 if attachment.disposition:
116 sfs["disposition"] = attachment.disposition
117 else:
118 del sfs["disposition"]
119 if thumbnail is not None:
120 sfs["file"].append(thumbnail)
121 if attachment.aio_stream is not None:
122 # revove the date in case we have downloaded the file to disk
123 # because no size was provided by the legacy module
124 sfs["file"]["date"] = None
125 if attachment.size:
126 sfs["file"]["size"] = attachment.size
127 stored.sfs = str(sfs)
128 msg.append(sfs)
130 def __send_url(
131 self,
132 msg: Message,
133 legacy_msg_id: str | None,
134 uploaded_url: str,
135 caption: str | None = None,
136 carbon: bool = False,
137 when: datetime | None = None,
138 correction: bool = False,
139 **kwargs: Any, # noqa:ANN401
140 ) -> list[Message]:
141 msg["oob"]["url"] = uploaded_url
142 msg["body"] = uploaded_url
143 if "sfs" in msg:
144 msg["fallback"].enable("body")
145 msg["fallback"]["for"] = self.xmpp.plugin["xep_0447"].stanza.NAMESPACE
146 if caption:
147 if correction:
148 if not legacy_msg_id:
149 raise TypeError
150 msg["replace"]["id"] = self._replace_id(legacy_msg_id)
151 elif legacy_msg_id:
152 self._set_msg_id(msg, legacy_msg_id)
153 m1 = self._send(msg, carbon=carbon, correction=correction, **kwargs)
154 m2 = self.send_text(
155 caption, legacy_msg_id=None, when=when, carbon=carbon, **kwargs
156 )
157 return [m1, m2] if m2 else [m1]
158 else:
159 if correction:
160 if not legacy_msg_id:
161 raise TypeError
162 msg["replace"]["id"] = self._replace_id(legacy_msg_id)
163 elif legacy_msg_id:
164 self._set_msg_id(msg, legacy_msg_id)
165 return [self._send(msg, carbon=carbon, **kwargs)]
167 def __get_base_message(
168 self,
169 legacy_msg_id: str | None = None,
170 reply_to: MessageReference | None = None,
171 when: datetime | None = None,
172 thread: str | None = None,
173 carbon: bool = False,
174 correction: bool = False,
175 mto: JID | None = None,
176 ) -> Message:
177 if correction:
178 if not legacy_msg_id:
179 raise TypeError
180 xmpp_ids = self._legacy_to_xmpp(legacy_msg_id)
181 if xmpp_ids:
182 original_xmpp_id = xmpp_ids[0]
183 for xmpp_id in xmpp_ids:
184 if xmpp_id == original_xmpp_id:
185 continue
186 self.retract(xmpp_id, thread)
188 if reply_to is not None and reply_to.body:
189 # We cannot have a "quote fallback" for attachments since most (all?)
190 # XMPP clients will only treat a message as an attachment if the
191 # body is the URL and nothing else.
192 reply_to_for_attachment: MessageReference | None = MessageReference(
193 reply_to.legacy_id, reply_to.author
194 )
195 else:
196 reply_to_for_attachment = reply_to
198 return self._make_message(
199 when=when,
200 reply_to=reply_to_for_attachment,
201 carbon=carbon,
202 mto=mto,
203 thread=thread,
204 )
206 async def send_file(
207 self,
208 attachment: LegacyAttachment | Path | str,
209 legacy_msg_id: str | None = None,
210 *,
211 reply_to: MessageReference | None = None,
212 when: datetime | None = None,
213 thread: str | None = None,
214 **kwargs: Any, # noqa:ANN401
215 ) -> tuple[str | None, list[Message]]:
216 """
217 Send a single file from this :term:`XMPP Entity`.
219 :param attachment: The file to send.
220 Ideally, a :class:`.LegacyAttachment` with a unique ``legacy_file_id``
221 attribute set, to optimise potential future reuses.
222 It can also be:
223 - a :class:`pathlib.Path` instance to point to a local file, or
224 - a ``str``, representing a fetchable HTTP URL.
225 :param legacy_msg_id: If you want to be able to transport read markers from the gateway
226 user to the legacy network, specify this
227 :param reply_to: Quote another message (:xep:`0461`)
228 :param when: when the file was sent, for a "delay" tag (:xep:`0203`)
229 :param thread:
230 """
231 async with self._uploader.dedup_lock(attachment):
232 return await self.__send_file(
233 attachment,
234 legacy_msg_id,
235 reply_to=reply_to,
236 when=when,
237 thread=thread,
238 **kwargs,
239 )
241 async def __send_file(
242 self,
243 attachment: LegacyAttachment | Path | str,
244 legacy_msg_id: str | None = None,
245 *,
246 reply_to: MessageReference | None = None,
247 when: datetime | None = None,
248 thread: str | None = None,
249 store_multi: bool = True,
250 carbon: bool = False,
251 mto: JID | None = None,
252 correction: bool = False,
253 **send_kwargs: Any, # noqa:ANN401
254 ) -> tuple[str | None, list[Message]]:
255 if isinstance(attachment, str):
256 attachment = LegacyAttachment(url=attachment)
257 elif isinstance(attachment, Path):
258 attachment = LegacyAttachment(path=attachment)
260 msg = self.__get_base_message(
261 legacy_msg_id=legacy_msg_id,
262 reply_to=reply_to,
263 when=when,
264 thread=thread,
265 carbon=carbon,
266 correction=correction,
267 mto=mto,
268 )
269 if attachment.is_sticker:
270 msg.enable("sticker")
272 stored = await self._uploader.get_stored(attachment)
274 if attachment.content_type is None and (
275 name := (attachment.name or attachment.url or attachment.path)
276 ):
277 attachment.content_type, _ = guess_type(name)
279 if not attachment.name:
280 if attachment.url:
281 attachment.url.split("/")[-1]
282 elif isinstance(attachment.path, Path):
283 attachment.name = attachment.path.name
284 else:
285 attachment.name = uuid4().hex
286 if attachment.content_type:
287 ext = guess_extension(attachment.content_type)
288 if ext:
289 attachment.name += ext
291 try:
292 new_url = (
293 stored.url
294 if stored.url
295 else await self._uploader.get_url(attachment, stored)
296 )
297 except Exception as e:
298 log.error("Error with attachment: %s: %s", attachment, e)
299 log.debug("", exc_info=e)
300 msg["body"] = (
301 f"/me tried to send a file ({attachment.format_for_user()}), "
302 f"but something went wrong: {e}. "
303 )
304 self._set_msg_id(msg, legacy_msg_id)
305 return None, [self._send(msg, **send_kwargs)]
306 else:
307 attachment.url = new_url
308 thumbnail = await self.__get_thumbnail(attachment)
309 self.__set_sims(msg, new_url, attachment, stored, thumbnail)
310 self.__set_sfs(msg, new_url, attachment, stored, thumbnail)
311 finally:
312 if isinstance(attachment.path, Path) and is_temp_path(
313 attachment.path, async_iterator_download_only=True
314 ):
315 try:
316 attachment.path.unlink()
317 attachment.path.parent.rmdir()
318 except (OSError, FileNotFoundError):
319 log.exception("Failed cleaning up %s", attachment.path)
320 attachment.path = None
321 stored.url = new_url
323 self._uploader.record(stored)
325 msgs = self.__send_url(
326 msg,
327 legacy_msg_id,
328 uploaded_url=new_url,
329 caption=attachment.caption,
330 carbon=carbon,
331 when=when,
332 correction=correction,
333 **send_kwargs,
334 )
335 if store_multi and legacy_msg_id:
336 self.__store_multi(legacy_msg_id, msgs)
338 return new_url, msgs
340 def __send_body(
341 self,
342 body: str | None = None,
343 legacy_msg_id: str | None = None,
344 reply_to: MessageReference | None = None,
345 when: datetime | None = None,
346 thread: str | None = None,
347 **kwargs: Any, # noqa:ANN401
348 ) -> Message | None:
349 if body:
350 return self.send_text(
351 body,
352 legacy_msg_id,
353 reply_to=reply_to,
354 when=when,
355 thread=thread,
356 **kwargs,
357 )
358 else:
359 return None
361 async def send_files(
362 self,
363 attachments: Collection[LegacyAttachment],
364 legacy_msg_id: str | None = None,
365 body: str | None = None,
366 *,
367 reply_to: MessageReference | None = None,
368 when: datetime | None = None,
369 thread: str | None = None,
370 body_first: bool = False,
371 correction: bool = False,
372 correction_event_id: str | None = None,
373 **kwargs: Any, # noqa:ANN401
374 ) -> None:
375 # TODO: once the epic XEP-0385 vs XEP-0447 battle is over, pick
376 # one and stop sending several attachments this way
377 # we attach the legacy_message ID to the last message we send, because
378 # we don't want several messages with the same ID (especially for MUC MAM)
379 if not attachments and not body:
380 # ignoring empty message
381 return
382 body_msg_id = (
383 legacy_msg_id if body_needs_msg_id(attachments, body, body_first) else None
384 )
385 send_body = functools.partial(
386 self.__send_body,
387 body=body,
388 reply_to=reply_to,
389 when=when,
390 thread=thread,
391 correction=correction,
392 legacy_msg_id=body_msg_id,
393 correction_event_id=correction_event_id,
394 **kwargs,
395 )
396 all_msgs = []
397 if body_first:
398 all_msgs.append(send_body())
399 for i, attachment in enumerate(attachments):
400 legacy = legacy_msg_id if i == 0 and body_msg_id is None else None
401 _url, msgs = await self.send_file(
402 attachment,
403 legacy,
404 reply_to=reply_to,
405 when=when,
406 thread=thread,
407 store_multi=False,
408 **kwargs,
409 )
410 all_msgs.extend(msgs)
411 if not body_first:
412 all_msgs.append(send_body())
413 self.__store_multi(legacy_msg_id, all_msgs)
415 def __store_multi(
416 self,
417 legacy_msg_id: str | None,
418 all_msgs: Sequence[Message | None],
419 ) -> None:
420 if legacy_msg_id is None:
421 return
422 if (recipient_pk := self._recipient_pk()) is None:
423 return
424 ids = []
425 for msg in all_msgs:
426 if not msg:
427 continue
428 if stanza_id := msg.get_plugin("stanza_id", check=True):
429 ids.append(stanza_id["id"])
430 else:
431 ids.append(msg.get_id())
432 with self.xmpp.store.session() as orm:
433 self.xmpp.store.id_map.set_msg(
434 orm, recipient_pk, str(legacy_msg_id), ids, self.is_participant
435 )
436 orm.commit()
438 async def __get_thumbnail(self, attachment: LegacyAttachment) -> Thumbnail | None:
439 if attachment.content_type is None:
440 return None
441 if not (attachment.data or attachment.stream or attachment.path):
442 return None
443 if not attachment.content_type.startswith("image"):
444 return None
446 try:
447 h, x, y = await self.xmpp.loop.run_in_executor(
448 avatar_cache._thread_pool,
449 get_thumbhash,
450 attachment,
451 )
452 except Exception as e:
453 log.debug("Could not generate a thumbhash", exc_info=e)
454 return None
456 thumbnail = Thumbnail()
457 thumbnail["width"] = x
458 thumbnail["height"] = y
459 thumbnail["media-type"] = "image/thumbhash"
460 thumbnail["uri"] = "data:image/thumbhash;base64," + urlquote(h)
461 return thumbnail
464def body_needs_msg_id(
465 attachments: Collection[LegacyAttachment], body: str | None, body_first: bool
466) -> bool:
467 if attachments:
468 return bool(body and body_first)
469 else:
470 return True
473def get_thumbhash(att: LegacyAttachment) -> tuple[str, int, int]:
474 img, width, height = get_image(att)
475 rgba_2d = img.get_flattened_data()
476 assert isinstance(rgba_2d[0], tuple)
477 rgba = list(chain(*rgba_2d))
478 ints = thumbhash.rgba_to_thumb_hash(img.width, img.height, rgba)
479 return base64.b64encode(bytes(ints)).decode(), width, height
482def get_image(att: LegacyAttachment) -> tuple[Image.Image, int, int]:
483 if att.data:
484 img = Image.open(io.BytesIO(att.data))
485 elif att.path:
486 with cast(Path, att.path).open("rb") as fp:
487 img = Image.open(fp)
488 img.load()
489 elif att.stream:
490 att.stream.seek(0)
491 img = Image.open(att.stream)
492 else:
493 raise RuntimeError("No way to read the image")
494 width, height = img.size
495 rgba = img.convert("RGBA")
496 if width > 100 or height > 100:
497 rgba.thumbnail((100, 100))
498 return ImageOps.exif_transpose(rgba), width, height
501log = logging.getLogger(__name__)