Coverage for slidge/core/mixins/attachment.py: 89%
197 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-29 05:05 +0000
1from __future__ import annotations
3import base64
4import functools
5import io
6import logging
7import warnings
8from collections.abc import Collection, Sequence
9from datetime import datetime
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
16try:
17 import thash
18except ImportError:
19 warnings.warn(
20 "thash is not installed, slidge won't be able to generate "
21 "thumbhashes for image attachments."
22 )
23 _THUMBHASH_IMPORTED = False
24else:
25 _THUMBHASH_IMPORTED = True
27from PIL import Image
28from slixmpp import JID, Message
30from ...db.avatar import avatar_cache
31from ...util.types import LegacyAttachment, MessageReference
32from ..attachment_upload import is_temp_path
33from .message_text import TextMessageMixin
36class AttachmentMixin(TextMessageMixin):
37 def __send_url(
38 self,
39 msg: Message,
40 legacy_msg_id: str | None,
41 uploaded_url: str,
42 caption: str | None = None,
43 carbon: bool = False,
44 when: datetime | None = None,
45 correction: bool = False,
46 **kwargs: Any, # noqa:ANN401
47 ) -> list[Message]:
48 msg["oob"]["url"] = uploaded_url
49 msg["body"] = uploaded_url
50 if "sfs" in msg:
51 msg["fallback"].enable("body")
52 msg["fallback"]["for"] = self.xmpp.plugin["xep_0447"].stanza.NAMESPACE
53 if caption:
54 if correction:
55 if not legacy_msg_id:
56 raise TypeError
57 msg["replace"]["id"] = self._replace_id(legacy_msg_id)
58 elif legacy_msg_id:
59 self._set_msg_id(msg, legacy_msg_id)
60 m1 = self._send(msg, carbon=carbon, correction=correction, **kwargs)
61 m2 = self.send_text(
62 caption, legacy_msg_id=None, when=when, carbon=carbon, **kwargs
63 )
64 return [m1, m2] if m2 else [m1]
65 else:
66 if correction:
67 if not legacy_msg_id:
68 raise TypeError
69 msg["replace"]["id"] = self._replace_id(legacy_msg_id)
70 elif legacy_msg_id:
71 self._set_msg_id(msg, legacy_msg_id)
72 return [self._send(msg, carbon=carbon, **kwargs)]
74 def __get_base_message(
75 self,
76 legacy_msg_id: str | None = None,
77 reply_to: MessageReference | None = None,
78 when: datetime | None = None,
79 thread: str | None = None,
80 carbon: bool = False,
81 correction: bool = False,
82 mto: JID | None = None,
83 ) -> Message:
84 if correction:
85 if not legacy_msg_id:
86 raise TypeError
87 xmpp_ids = self._legacy_to_xmpp(legacy_msg_id)
88 if xmpp_ids:
89 original_xmpp_id = xmpp_ids[0]
90 for xmpp_id in xmpp_ids:
91 if xmpp_id == original_xmpp_id:
92 continue
93 self.retract(xmpp_id, thread)
95 if reply_to is not None and reply_to.body:
96 # We cannot have a "quote fallback" for attachments since most (all?)
97 # XMPP clients will only treat a message as an attachment if the
98 # body is the URL and nothing else.
99 reply_to_for_attachment: MessageReference | None = MessageReference(
100 reply_to.legacy_id, reply_to.author
101 )
102 else:
103 reply_to_for_attachment = reply_to
105 return self._make_message(
106 when=when,
107 reply_to=reply_to_for_attachment,
108 carbon=carbon,
109 mto=mto,
110 thread=thread,
111 )
113 async def send_file(
114 self,
115 attachment: LegacyAttachment | Path | str,
116 legacy_msg_id: str | None = None,
117 *,
118 reply_to: MessageReference | None = None,
119 when: datetime | None = None,
120 thread: str | None = None,
121 **kwargs: Any, # noqa:ANN401
122 ) -> tuple[str | None, list[Message]]:
123 """
124 Send a single file from this :term:`XMPP Entity`.
126 :param attachment: The file to send.
127 Ideally, a :class:`.LegacyAttachment` with a unique ``legacy_file_id``
128 attribute set, to optimise potential future reuses.
129 It can also be:
130 - a :class:`pathlib.Path` instance to point to a local file, or
131 - a ``str``, representing a fetchable HTTP URL.
132 :param legacy_msg_id: If you want to be able to transport read markers from the gateway
133 user to the legacy network, specify this
134 :param reply_to: Quote another message (:xep:`0461`)
135 :param when: when the file was sent, for a "delay" tag (:xep:`0203`)
136 :param thread:
137 """
138 async with self._uploader.dedup_lock(attachment):
139 return await self.__send_file(
140 attachment,
141 legacy_msg_id,
142 reply_to=reply_to,
143 when=when,
144 thread=thread,
145 **kwargs,
146 )
148 async def __send_file(
149 self,
150 attachment: LegacyAttachment | Path | str,
151 legacy_msg_id: str | None = None,
152 *,
153 reply_to: MessageReference | None = None,
154 when: datetime | None = None,
155 thread: str | None = None,
156 store_multi: bool = True,
157 carbon: bool = False,
158 mto: JID | None = None,
159 correction: bool = False,
160 **send_kwargs: Any, # noqa:ANN401
161 ) -> tuple[str | None, list[Message]]:
162 if isinstance(attachment, str):
163 attachment = LegacyAttachment(url=attachment)
164 elif isinstance(attachment, Path):
165 attachment = LegacyAttachment(path=attachment)
167 msg = self.__get_base_message(
168 legacy_msg_id=legacy_msg_id,
169 reply_to=reply_to,
170 when=when,
171 thread=thread,
172 carbon=carbon,
173 correction=correction,
174 mto=mto,
175 )
176 if attachment.is_sticker:
177 msg.enable("sticker")
179 stored = await self._uploader.get_stored(attachment)
181 if attachment.content_type is None and (
182 name := (attachment.name or attachment.url or attachment.path)
183 ):
184 attachment.content_type, _ = guess_type(name)
186 if not attachment.name:
187 if attachment.url:
188 attachment.url.split("/")[-1]
189 elif isinstance(attachment.path, Path):
190 attachment.name = attachment.path.name
191 else:
192 attachment.name = uuid4().hex
193 if attachment.content_type:
194 ext = guess_extension(attachment.content_type)
195 if ext:
196 attachment.name += ext
198 try:
199 new_url = (
200 stored.url
201 if stored.url
202 else await self._uploader.get_url(attachment, stored)
203 )
204 except Exception as e:
205 log.error("Error with attachment: %s: %s", attachment, e)
206 log.debug("", exc_info=e)
207 msg["body"] = (
208 f"/me tried to send a file ({attachment.format_for_user()}), "
209 f"but something went wrong: {e}. "
210 )
211 self._set_msg_id(msg, legacy_msg_id)
212 return None, [self._send(msg, **send_kwargs)]
213 else:
214 thumbhash_width_height = await self.__get_thumbhash(attachment)
215 if attachment.size is None:
216 if attachment.data:
217 attachment.size = len(attachment.data)
218 elif attachment.path:
219 assert isinstance(attachment.path, Path)
220 attachment.size = attachment.path.stat().st_size
221 finally:
222 if isinstance(attachment.path, Path) and is_temp_path(
223 attachment.path, async_iterator_download_only=True
224 ):
225 try:
226 attachment.path.unlink()
227 attachment.path.parent.rmdir()
228 except (OSError, FileNotFoundError):
229 log.exception("Failed cleaning up %s", attachment.path)
230 attachment.path = None
232 stored.url = attachment.url = new_url
233 if thumbhash_width_height is not None:
234 stored.thumbhash, stored.width, stored.height = thumbhash_width_height
236 stored.size = attachment.size
237 stored.content_type = attachment.content_type
238 stored.desc = attachment.caption
239 stored.name = attachment.name
240 stored.disposition = attachment.disposition
241 stored.date = attachment.date
242 stored.sha256 = attachment.sha256
244 msg.append(stored.sims())
245 msg.append(stored.sfs())
247 self._uploader.record(stored)
249 msgs = self.__send_url(
250 msg,
251 legacy_msg_id,
252 uploaded_url=new_url,
253 caption=attachment.caption,
254 carbon=carbon,
255 when=when,
256 correction=correction,
257 **send_kwargs,
258 )
259 if store_multi and legacy_msg_id:
260 self.__store_multi(legacy_msg_id, msgs)
262 return new_url, msgs
264 def __send_body(
265 self,
266 body: str | None = None,
267 legacy_msg_id: str | None = None,
268 reply_to: MessageReference | None = None,
269 when: datetime | None = None,
270 thread: str | None = None,
271 **kwargs: Any, # noqa:ANN401
272 ) -> Message | None:
273 if body:
274 return self.send_text(
275 body,
276 legacy_msg_id,
277 reply_to=reply_to,
278 when=when,
279 thread=thread,
280 **kwargs,
281 )
282 else:
283 return None
285 async def send_files(
286 self,
287 attachments: Collection[LegacyAttachment],
288 legacy_msg_id: str | None = None,
289 body: str | None = None,
290 *,
291 reply_to: MessageReference | None = None,
292 when: datetime | None = None,
293 thread: str | None = None,
294 body_first: bool = False,
295 correction: bool = False,
296 correction_event_id: str | None = None,
297 **kwargs: Any, # noqa:ANN401
298 ) -> None:
299 # TODO: once the epic XEP-0385 vs XEP-0447 battle is over, pick
300 # one and stop sending several attachments this way
301 # we attach the legacy_message ID to the last message we send, because
302 # we don't want several messages with the same ID (especially for MUC MAM)
303 if not attachments and not body:
304 # ignoring empty message
305 return
306 body_msg_id = (
307 legacy_msg_id if body_needs_msg_id(attachments, body, body_first) else None
308 )
309 send_body = functools.partial(
310 self.__send_body,
311 body=body,
312 reply_to=reply_to,
313 when=when,
314 thread=thread,
315 correction=correction,
316 legacy_msg_id=body_msg_id,
317 correction_event_id=correction_event_id,
318 **kwargs,
319 )
320 all_msgs = []
321 if body_first:
322 all_msgs.append(send_body())
323 for i, attachment in enumerate(attachments):
324 legacy = legacy_msg_id if i == 0 and body_msg_id is None else None
325 _url, msgs = await self.send_file(
326 attachment,
327 legacy,
328 reply_to=reply_to,
329 when=when,
330 thread=thread,
331 store_multi=False,
332 **kwargs,
333 )
334 all_msgs.extend(msgs)
335 if not body_first:
336 all_msgs.append(send_body())
337 self.__store_multi(legacy_msg_id, all_msgs)
339 def __store_multi(
340 self,
341 legacy_msg_id: str | None,
342 all_msgs: Sequence[Message | None],
343 ) -> None:
344 if legacy_msg_id is None:
345 return
346 if (recipient_pk := self._recipient_pk()) is None:
347 return
348 ids = []
349 for msg in all_msgs:
350 if not msg:
351 continue
352 if stanza_id := msg.get_plugin("stanza_id", check=True):
353 ids.append(stanza_id["id"])
354 else:
355 ids.append(msg.get_id())
356 with self.xmpp.store.session() as orm:
357 self.xmpp.store.id_map.set_msg(
358 orm, recipient_pk, str(legacy_msg_id), ids, self.is_participant
359 )
360 orm.commit()
362 async def __get_thumbhash(
363 self, attachment: LegacyAttachment
364 ) -> tuple[str, int, int] | None:
365 if not _THUMBHASH_IMPORTED:
366 return None
367 if attachment.content_type is None:
368 return None
369 if not (attachment.data or attachment.stream or attachment.path):
370 return None
371 if not attachment.content_type.startswith("image"):
372 return None
374 try:
375 h, x, y = await self.xmpp.loop.run_in_executor(
376 avatar_cache._thread_pool,
377 get_thumbhash,
378 attachment,
379 )
380 except Exception as e:
381 log.debug("Could not generate a thumbhash", exc_info=e)
382 return None
384 return urlquote(h), x, y
387def body_needs_msg_id(
388 attachments: Collection[LegacyAttachment], body: str | None, body_first: bool
389) -> bool:
390 if attachments:
391 return bool(body and body_first)
392 else:
393 return True
396def get_thumbhash(att: LegacyAttachment) -> tuple[str, int, int]:
397 img = get_image(att)
398 ints = thash.encode(img)
399 return base64.b64encode(bytes(ints)).decode(), img.width, img.height
402def get_image(att: LegacyAttachment) -> Image.Image:
403 if att.data:
404 return Image.open(io.BytesIO(att.data))
405 if att.path:
406 with cast(Path, att.path).open("rb") as fp:
407 img = Image.open(fp)
408 img.load()
409 return img
410 if att.stream:
411 att.stream.seek(0)
412 img = Image.open(att.stream)
414 raise RuntimeError("No way to read the image")
417log = logging.getLogger(__name__)