Coverage for slidge/core/mixins/attachment.py: 93%
246 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-18 04:30 +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 if stored.url:
292 new_url = stored.url
293 else:
294 try:
295 new_url = await self._uploader.get_url(attachment, stored)
296 except Exception as e:
297 log.error("Error with attachment: %s: %s", attachment, e)
298 log.debug("", exc_info=e)
299 msg["body"] = (
300 f"/me tried to send a file ({attachment.format_for_user()}), "
301 f"but something went wrong: {e}. "
302 )
303 self._set_msg_id(msg, legacy_msg_id)
304 return None, [self._send(msg, **send_kwargs)]
305 stored.url = new_url
306 attachment.url = new_url
307 thumbnail = await self.__get_thumbnail(attachment)
308 self.__set_sims(msg, new_url, attachment, stored, thumbnail)
309 self.__set_sfs(msg, new_url, attachment, stored, thumbnail)
311 if isinstance(attachment.path, Path) and is_temp_path(
312 attachment.path, async_iterator_download_only=True
313 ):
314 attachment.path.unlink()
315 attachment.path = None
317 self._uploader.record(stored)
319 msgs = self.__send_url(
320 msg,
321 legacy_msg_id,
322 uploaded_url=new_url,
323 caption=attachment.caption,
324 carbon=carbon,
325 when=when,
326 correction=correction,
327 **send_kwargs,
328 )
329 if store_multi and legacy_msg_id:
330 self.__store_multi(legacy_msg_id, msgs)
332 return new_url, msgs
334 def __send_body(
335 self,
336 body: str | None = None,
337 legacy_msg_id: str | None = None,
338 reply_to: MessageReference | None = None,
339 when: datetime | None = None,
340 thread: str | None = None,
341 **kwargs: Any, # noqa:ANN401
342 ) -> Message | None:
343 if body:
344 return self.send_text(
345 body,
346 legacy_msg_id,
347 reply_to=reply_to,
348 when=when,
349 thread=thread,
350 **kwargs,
351 )
352 else:
353 return None
355 async def send_files(
356 self,
357 attachments: Collection[LegacyAttachment],
358 legacy_msg_id: str | None = None,
359 body: str | None = None,
360 *,
361 reply_to: MessageReference | None = None,
362 when: datetime | None = None,
363 thread: str | None = None,
364 body_first: bool = False,
365 correction: bool = False,
366 correction_event_id: str | None = None,
367 **kwargs: Any, # noqa:ANN401
368 ) -> None:
369 # TODO: once the epic XEP-0385 vs XEP-0447 battle is over, pick
370 # one and stop sending several attachments this way
371 # we attach the legacy_message ID to the last message we send, because
372 # we don't want several messages with the same ID (especially for MUC MAM)
373 if not attachments and not body:
374 # ignoring empty message
375 return
376 body_msg_id = (
377 legacy_msg_id if body_needs_msg_id(attachments, body, body_first) else None
378 )
379 send_body = functools.partial(
380 self.__send_body,
381 body=body,
382 reply_to=reply_to,
383 when=when,
384 thread=thread,
385 correction=correction,
386 legacy_msg_id=body_msg_id,
387 correction_event_id=correction_event_id,
388 **kwargs,
389 )
390 all_msgs = []
391 if body_first:
392 all_msgs.append(send_body())
393 for i, attachment in enumerate(attachments):
394 legacy = legacy_msg_id if i == 0 and body_msg_id is None else None
395 _url, msgs = await self.send_file(
396 attachment,
397 legacy,
398 reply_to=reply_to,
399 when=when,
400 thread=thread,
401 store_multi=False,
402 **kwargs,
403 )
404 all_msgs.extend(msgs)
405 if not body_first:
406 all_msgs.append(send_body())
407 self.__store_multi(legacy_msg_id, all_msgs)
409 def __store_multi(
410 self,
411 legacy_msg_id: str | None,
412 all_msgs: Sequence[Message | None],
413 ) -> None:
414 if legacy_msg_id is None:
415 return
416 if (recipient_pk := self._recipient_pk()) is None:
417 return
418 ids = []
419 for msg in all_msgs:
420 if not msg:
421 continue
422 if stanza_id := msg.get_plugin("stanza_id", check=True):
423 ids.append(stanza_id["id"])
424 else:
425 ids.append(msg.get_id())
426 with self.xmpp.store.session() as orm:
427 self.xmpp.store.id_map.set_msg(
428 orm, recipient_pk, str(legacy_msg_id), ids, self.is_participant
429 )
430 orm.commit()
432 async def __get_thumbnail(self, attachment: LegacyAttachment) -> Thumbnail | None:
433 if attachment.content_type is None:
434 return None
435 if not (attachment.data or attachment.stream or attachment.path):
436 return None
437 if not attachment.content_type.startswith("image"):
438 return None
440 try:
441 h, x, y = await self.xmpp.loop.run_in_executor(
442 avatar_cache._thread_pool,
443 get_thumbhash,
444 attachment,
445 )
446 except Exception as e:
447 log.debug("Could not generate a thumbhash", exc_info=e)
448 return None
450 thumbnail = Thumbnail()
451 thumbnail["width"] = x
452 thumbnail["height"] = y
453 thumbnail["media-type"] = "image/thumbhash"
454 thumbnail["uri"] = "data:image/thumbhash;base64," + urlquote(h)
455 return thumbnail
458def body_needs_msg_id(
459 attachments: Collection[LegacyAttachment], body: str | None, body_first: bool
460) -> bool:
461 if attachments:
462 return bool(body and body_first)
463 else:
464 return True
467def get_thumbhash(att: LegacyAttachment) -> tuple[str, int, int]:
468 img, width, height = get_image(att)
469 rgba_2d = img.get_flattened_data()
470 assert isinstance(rgba_2d[0], tuple)
471 rgba = list(chain(*rgba_2d))
472 ints = thumbhash.rgba_to_thumb_hash(img.width, img.height, rgba)
473 return base64.b64encode(bytes(ints)).decode(), width, height
476def get_image(att: LegacyAttachment) -> tuple[Image.Image, int, int]:
477 if att.data:
478 img = Image.open(io.BytesIO(att.data))
479 elif att.path:
480 with cast(Path, att.path).open("rb") as fp:
481 img = Image.open(fp)
482 img.load()
483 elif att.stream:
484 att.stream.seek(0)
485 img = Image.open(att.stream)
486 else:
487 raise RuntimeError("No way to read the image")
488 width, height = img.size
489 rgba = img.convert("RGBA")
490 if width > 100 or height > 100:
491 rgba.thumbnail((100, 100))
492 return ImageOps.exif_transpose(rgba), width, height
495log = logging.getLogger(__name__)