Coverage for slidge/core/mixins/attachment.py: 89%

424 statements  

« prev     ^ index     » next       coverage.py v7.15.0, created at 2026-07-09 03:59 +0000

1from __future__ import annotations 

2 

3import base64 

4import contextlib 

5import functools 

6import io 

7import logging 

8import os 

9import shutil 

10import stat 

11import tempfile 

12import warnings 

13from collections.abc import AsyncIterator, Collection, Sequence 

14from datetime import datetime 

15from itertools import chain 

16from mimetypes import guess_extension, guess_type 

17from pathlib import Path 

18from typing import IO, Any, cast 

19from urllib.parse import quote as urlquote 

20from uuid import uuid4 

21from xml.etree import ElementTree as ET 

22 

23import aiohttp 

24import thumbhash 

25from PIL import Image, ImageOps 

26from slixmpp import JID, Iq, Message 

27from slixmpp.plugins.xep_0264.stanza import Thumbnail 

28from slixmpp.plugins.xep_0447.stanza import StatelessFileSharing 

29 

30from ...db.avatar import avatar_cache 

31from ...db.models import Attachment 

32from ...util.types import LegacyAttachment, MessageReference 

33from ...util.util import fix_suffix 

34from .. import config 

35from .message_text import TextMessageMixin 

36 

37 

38class AttachmentMixin(TextMessageMixin): 

39 @property 

40 def __is_component(self) -> bool: 

41 return self.session is NotImplemented 

42 

43 async def __upload(self, att: _AttachmentWithName) -> str: 

44 assert config.UPLOAD_SERVICE 

45 att = await _ensure_metadata(att) 

46 iq_slot = await self.__request_upload_slot( 

47 config.UPLOAD_SERVICE, 

48 att.name, 

49 att.size, 

50 att.content_type, 

51 ) 

52 slot = iq_slot["http_upload_slot"] 

53 headers = { 

54 "Content-Length": str(att.size), 

55 "Content-Type": att.content_type, 

56 **{header["name"]: header["value"] for header in slot["put"]["headers"]}, 

57 } 

58 

59 async with ( 

60 aiohttp.ClientSession() as http, 

61 _get_data(att, http) as data, 

62 http.put(slot["put"]["url"], data=data, headers=headers) as resp, 

63 ): 

64 resp.raise_for_status() 

65 

66 return slot["get"]["url"] # type:ignore[no-any-return] 

67 

68 async def __request_upload_slot( 

69 self, 

70 upload_service: JID | str, 

71 filename: str, 

72 size: int, 

73 content_type: str, 

74 ) -> Iq: 

75 iq_request = self.xmpp.make_iq_get(ito=upload_service) 

76 request = iq_request["http_upload_request"] 

77 request["filename"] = filename 

78 request["size"] = str(size) 

79 request["content-type"] = content_type 

80 if not self.__is_component: 

81 iq_request.set_from(self.session.user_jid) 

82 try: 

83 return await self.xmpp.plugin["xep_0356"].send_privileged_iq(iq_request) 

84 except Exception as e: 

85 warnings.warn( 

86 "Could not request upload slot on behalf of " 

87 f"{self.session.user_jid}: {e}." 

88 "Falling back to not using privileges." 

89 ) 

90 iq_request.set_from(config.UPLOAD_REQUESTER or self.xmpp.boundjid) 

91 return await iq_request.send() # type:ignore[no-any-return] 

92 

93 @staticmethod 

94 async def __no_upload( 

95 att: _AttachmentWithName, legacy_file_id: str | None 

96 ) -> tuple[Path, str]: 

97 file_id = uuid4().hex if legacy_file_id is None else legacy_file_id 

98 assert config.NO_UPLOAD_PATH is not None 

99 assert config.NO_UPLOAD_URL_PREFIX is not None 

100 destination_dir = Path(config.NO_UPLOAD_PATH) / file_id 

101 

102 if destination_dir.exists(): 

103 log.debug("Dest dir exists: %s", destination_dir) 

104 files = list(f for f in destination_dir.glob("**/*") if f.is_file()) 

105 if len(files) == 1: 

106 log.debug( 

107 "Found the legacy attachment '%s' at '%s'", 

108 legacy_file_id, 

109 files[0], 

110 ) 

111 name = files[0].name 

112 uu = files[0].parent.name # anti-obvious url trick, see below 

113 return files[0], "/".join([file_id, uu, name]) 

114 else: 

115 log.warning( 

116 ( 

117 "There are several or zero files in %s, " 

118 "slidge doesn't know which one to pick among %s. " 

119 "Removing the dir." 

120 ), 

121 destination_dir, 

122 files, 

123 ) 

124 shutil.rmtree(destination_dir) 

125 

126 log.debug("Did not find a file in: %s", destination_dir) 

127 # let's use a UUID to avoid URLs being too obvious 

128 uu = str(uuid4()) 

129 destination_dir = destination_dir / uu 

130 destination_dir.mkdir(parents=True) 

131 

132 assert att.name 

133 destination = destination_dir / att.name 

134 if att.path: 

135 assert isinstance(att.path, Path) 

136 try: 

137 destination.hardlink_to(att.path) 

138 except OSError as e: 

139 if is_temp_path(att.path): 

140 shutil.copy2(att.path, destination) 

141 else: 

142 log.debug("Could not hardlink: %s, attempting symlink", e) 

143 try: 

144 destination.symlink_to(att.path) 

145 except OSError as e: 

146 log.debug("Could not symlink: %s, copying data", e) 

147 shutil.copy2(att.path, destination) 

148 elif att.data: 

149 destination.write_bytes(att.data) 

150 else: 

151 with destination.open("wb") as f: 

152 if att.aio_stream: 

153 async for chunk in att.aio_stream: 

154 f.write(chunk) 

155 elif att.stream: 

156 shutil.copyfileobj(att.stream, f) 

157 elif att.url: 

158 async with ( 

159 aiohttp.ClientSession() as http, 

160 http.get(att.url) as resp, 

161 ): 

162 resp.raise_for_status() 

163 async for chunk in resp.content.iter_chunked(64 * 1024): 

164 f.write(chunk) 

165 else: 

166 raise RuntimeError 

167 

168 if config.NO_UPLOAD_FILE_READ_OTHERS: 

169 log.debug("Changing perms of %s", destination) 

170 destination.chmod(destination.stat().st_mode | stat.S_IROTH) 

171 

172 url = "/".join([file_id, uu, att.name]) 

173 return destination, url 

174 

175 async def __valid_url(self, url: str) -> bool: 

176 async with self.session.http.head(url) as r: 

177 return r.status < 400 

178 

179 async def __get_stored(self, attachment: LegacyAttachment) -> Attachment: 

180 if attachment.legacy_file_id is not None and not self.__is_component: 

181 with self.xmpp.store.session() as orm: 

182 stored = ( 

183 orm.query(Attachment) 

184 .filter_by( 

185 legacy_file_id=str(attachment.legacy_file_id), 

186 user_account_id=self.session.user_pk, 

187 ) 

188 .one_or_none() 

189 ) 

190 if stored is not None: 

191 if not await self.__valid_url(stored.url): 

192 stored.url = None # type:ignore 

193 return stored 

194 return Attachment( 

195 user_account_id=None if self.__is_component else self.session.user_pk, 

196 legacy_file_id=None 

197 if attachment.legacy_file_id is None 

198 else str(attachment.legacy_file_id), 

199 url=attachment.url if config.USE_ATTACHMENT_ORIGINAL_URLS else None, 

200 ) 

201 

202 async def __get_url(self, att: LegacyAttachment, stored: Attachment) -> str: 

203 att = _ensure_name(att) 

204 

205 if len(att.name) > config.ATTACHMENT_MAXIMUM_FILE_NAME_LENGTH: 

206 log.debug("Trimming long filename: %s", att.name) 

207 base, ext = os.path.splitext(att.name) 

208 att.name = ( 

209 base[: config.ATTACHMENT_MAXIMUM_FILE_NAME_LENGTH - len(ext)] + ext 

210 ) 

211 

212 if config.FIX_FILENAME_SUFFIX_MIME_TYPE and isinstance(att.path, Path): 

213 att.name, att.content_type = fix_suffix( 

214 att.path, att.content_type, att.name 

215 ) 

216 

217 att.legacy_file_id = stored.legacy_file_id 

218 

219 if config.NO_UPLOAD_PATH: 

220 att.path, new_url = await self.__no_upload(att, stored.legacy_file_id) 

221 new_url = (config.NO_UPLOAD_URL_PREFIX or "") + "/" + urlquote(new_url) 

222 else: 

223 new_url = await self.__upload(att) 

224 

225 if stored.legacy_file_id: 

226 stored.url = new_url 

227 

228 return new_url 

229 

230 def __set_sims( 

231 self, 

232 msg: Message, 

233 uploaded_url: str, 

234 attachment: LegacyAttachment, 

235 stored: Attachment, 

236 thumbnail: Thumbnail | None, 

237 ) -> None: 

238 if stored.sims is not None: 

239 ref = self.xmpp.plugin["xep_0372"].stanza.Reference( 

240 xml=ET.fromstring(stored.sims) 

241 ) 

242 msg.append(ref) 

243 return 

244 

245 if attachment.data or attachment.stream or attachment.path: 

246 ref = self.xmpp.plugin["xep_0385"].get_sims( 

247 attachment.path, # type:ignore[arg-type] 

248 [uploaded_url], 

249 attachment.content_type, 

250 attachment.caption, 

251 file=attachment.stream, 

252 data=attachment.data, 

253 ) 

254 else: 

255 sims = self.xmpp.plugin["xep_0385"].stanza.Sims() 

256 ref = self.xmpp.plugin["xep_0372"].stanza.Reference() 

257 ref["uri"] = uploaded_url 

258 ref["type"] = "data" 

259 sims["sources"].append(ref) 

260 ref = self.xmpp.plugin["xep_0372"].stanza.Reference() 

261 ref.append(sims) 

262 ref["type"] = "data" 

263 sims.enable("file") 

264 if attachment.content_type: 

265 sims["file"]["media-type"] = attachment.content_type 

266 if attachment.caption: 

267 sims["file"]["desc"] = attachment.caption 

268 if attachment.size: 

269 sims["file"]["size"] = attachment.size 

270 if attachment.name: 

271 ref["sims"]["file"]["name"] = attachment.name 

272 if attachment.aio_stream is not None: 

273 # revove the date in case we have downloaded the file to disk 

274 # because no size was provided by the legacy module 

275 ref["sims"]["file"]["date"] = None 

276 

277 if thumbnail is not None: 

278 ref["sims"]["file"].append(thumbnail) 

279 

280 stored.sims = str(ref) 

281 msg.append(ref) 

282 

283 def __set_sfs( 

284 self, 

285 msg: Message, 

286 uploaded_url: str, 

287 attachment: LegacyAttachment, 

288 stored: Attachment, 

289 thumbnail: Thumbnail | None = None, 

290 ) -> None: 

291 if stored.sfs is not None: 

292 msg.append(StatelessFileSharing(xml=ET.fromstring(stored.sfs))) 

293 return 

294 

295 if attachment.path or attachment.data or attachment.stream: 

296 sfs = self.xmpp.plugin["xep_0447"].get_sfs( # type:ignore 

297 attachment.path, 

298 [uploaded_url], 

299 attachment.content_type, 

300 attachment.caption, 

301 data=attachment.data, 

302 file=attachment.stream, 

303 ) 

304 else: 

305 sfs = self.xmpp.plugin["xep_0447"].stanza.StatelessFileSharing() 

306 ref = self.xmpp.plugin["xep_0447"].stanza.UrlData() 

307 ref["target"] = uploaded_url 

308 sfs["sources"].append(ref) 

309 sfs.enable("file") 

310 if attachment.content_type: 

311 sfs["file"]["media-type"] = attachment.content_type 

312 if attachment.name: 

313 sfs["file"]["name"] = attachment.name 

314 if attachment.disposition: 

315 sfs["disposition"] = attachment.disposition 

316 else: 

317 del sfs["disposition"] 

318 if thumbnail is not None: 

319 sfs["file"].append(thumbnail) 

320 if attachment.aio_stream is not None: 

321 # revove the date in case we have downloaded the file to disk 

322 # because no size was provided by the legacy module 

323 sfs["file"]["date"] = None 

324 if attachment.size: 

325 sfs["file"]["size"] = attachment.size 

326 stored.sfs = str(sfs) 

327 msg.append(sfs) 

328 

329 def __send_url( 

330 self, 

331 msg: Message, 

332 legacy_msg_id: str | None, 

333 uploaded_url: str, 

334 caption: str | None = None, 

335 carbon: bool = False, 

336 when: datetime | None = None, 

337 correction: bool = False, 

338 **kwargs: Any, # noqa:ANN401 

339 ) -> list[Message]: 

340 msg["oob"]["url"] = uploaded_url 

341 msg["body"] = uploaded_url 

342 if "sfs" in msg: 

343 msg["fallback"].enable("body") 

344 msg["fallback"]["for"] = self.xmpp.plugin["xep_0447"].stanza.NAMESPACE 

345 if caption: 

346 if correction: 

347 if not legacy_msg_id: 

348 raise TypeError 

349 msg["replace"]["id"] = self._replace_id(legacy_msg_id) 

350 elif legacy_msg_id: 

351 self._set_msg_id(msg, legacy_msg_id) 

352 m1 = self._send(msg, carbon=carbon, correction=correction, **kwargs) 

353 m2 = self.send_text( 

354 caption, legacy_msg_id=None, when=when, carbon=carbon, **kwargs 

355 ) 

356 return [m1, m2] if m2 else [m1] 

357 else: 

358 if correction: 

359 if not legacy_msg_id: 

360 raise TypeError 

361 msg["replace"]["id"] = self._replace_id(legacy_msg_id) 

362 elif legacy_msg_id: 

363 self._set_msg_id(msg, legacy_msg_id) 

364 return [self._send(msg, carbon=carbon, **kwargs)] 

365 

366 def __get_base_message( 

367 self, 

368 legacy_msg_id: str | None = None, 

369 reply_to: MessageReference | None = None, 

370 when: datetime | None = None, 

371 thread: str | None = None, 

372 carbon: bool = False, 

373 correction: bool = False, 

374 mto: JID | None = None, 

375 ) -> Message: 

376 if correction: 

377 if not legacy_msg_id: 

378 raise TypeError 

379 xmpp_ids = self._legacy_to_xmpp(legacy_msg_id) 

380 if xmpp_ids: 

381 original_xmpp_id = xmpp_ids[0] 

382 for xmpp_id in xmpp_ids: 

383 if xmpp_id == original_xmpp_id: 

384 continue 

385 self.retract(xmpp_id, thread) 

386 

387 if reply_to is not None and reply_to.body: 

388 # We cannot have a "quote fallback" for attachments since most (all?) 

389 # XMPP clients will only treat a message as an attachment if the 

390 # body is the URL and nothing else. 

391 reply_to_for_attachment: MessageReference | None = MessageReference( 

392 reply_to.legacy_id, reply_to.author 

393 ) 

394 else: 

395 reply_to_for_attachment = reply_to 

396 

397 return self._make_message( 

398 when=when, 

399 reply_to=reply_to_for_attachment, 

400 carbon=carbon, 

401 mto=mto, 

402 thread=thread, 

403 ) 

404 

405 async def send_file( 

406 self, 

407 attachment: LegacyAttachment | Path | str, 

408 legacy_msg_id: str | None = None, 

409 *, 

410 reply_to: MessageReference | None = None, 

411 when: datetime | None = None, 

412 thread: str | None = None, 

413 **kwargs: Any, # noqa:ANN401 

414 ) -> tuple[str | None, list[Message]]: 

415 """ 

416 Send a single file from this :term:`XMPP Entity`. 

417 

418 :param attachment: The file to send. 

419 Ideally, a :class:`.LegacyAttachment` with a unique ``legacy_file_id`` 

420 attribute set, to optimise potential future reuses. 

421 It can also be: 

422 - a :class:`pathlib.Path` instance to point to a local file, or 

423 - a ``str``, representing a fetchable HTTP URL. 

424 :param legacy_msg_id: If you want to be able to transport read markers from the gateway 

425 user to the legacy network, specify this 

426 :param reply_to: Quote another message (:xep:`0461`) 

427 :param when: when the file was sent, for a "delay" tag (:xep:`0203`) 

428 :param thread: 

429 """ 

430 coro = self.__send_file( 

431 attachment, 

432 legacy_msg_id, 

433 reply_to=reply_to, 

434 when=when, 

435 thread=thread, 

436 **kwargs, 

437 ) 

438 if ( 

439 self.__is_component 

440 or not isinstance(attachment, LegacyAttachment) 

441 or attachment.legacy_file_id is None 

442 ): 

443 return await coro 

444 else: 

445 # prevents race conditions where we download the same thing several time 

446 # and end up attempting to insert it twice in the DB, raising an 

447 # IntegrityError. 

448 async with self.session.lock(("attachment", attachment.legacy_file_id)): 

449 return await coro 

450 

451 async def __send_file( 

452 self, 

453 attachment: LegacyAttachment | Path | str, 

454 legacy_msg_id: str | None = None, 

455 *, 

456 reply_to: MessageReference | None = None, 

457 when: datetime | None = None, 

458 thread: str | None = None, 

459 store_multi: bool = True, 

460 carbon: bool = False, 

461 mto: JID | None = None, 

462 correction: bool = False, 

463 **send_kwargs: Any, # noqa:ANN401 

464 ) -> tuple[str | None, list[Message]]: 

465 if isinstance(attachment, str): 

466 attachment = LegacyAttachment(url=attachment) 

467 elif isinstance(attachment, Path): 

468 attachment = LegacyAttachment(path=attachment) 

469 

470 msg = self.__get_base_message( 

471 legacy_msg_id=legacy_msg_id, 

472 reply_to=reply_to, 

473 when=when, 

474 thread=thread, 

475 carbon=carbon, 

476 correction=correction, 

477 mto=mto, 

478 ) 

479 if attachment.is_sticker: 

480 msg.enable("sticker") 

481 

482 stored = await self.__get_stored(attachment) 

483 

484 if attachment.content_type is None and ( 

485 name := (attachment.name or attachment.url or attachment.path) 

486 ): 

487 attachment.content_type, _ = guess_type(name) 

488 

489 if not attachment.name: 

490 if attachment.url: 

491 attachment.url.split("/")[-1] 

492 elif isinstance(attachment.path, Path): 

493 attachment.name = attachment.path.name 

494 else: 

495 attachment.name = uuid4().hex 

496 if attachment.content_type: 

497 ext = guess_extension(attachment.content_type) 

498 if ext: 

499 attachment.name += ext 

500 

501 if stored.url: 

502 new_url = stored.url 

503 else: 

504 try: 

505 new_url = await self.__get_url(attachment, stored) 

506 except Exception as e: 

507 log.error("Error with attachment: %s: %s", attachment, e) 

508 log.debug("", exc_info=e) 

509 msg["body"] = ( 

510 f"/me tried to send a file ({attachment.format_for_user()}), " 

511 f"but something went wrong: {e}. " 

512 ) 

513 self._set_msg_id(msg, legacy_msg_id) 

514 return None, [self._send(msg, **send_kwargs)] 

515 stored.url = new_url 

516 attachment.url = new_url 

517 thumbnail = await self.__get_thumbnail(attachment) 

518 self.__set_sims(msg, new_url, attachment, stored, thumbnail) 

519 self.__set_sfs(msg, new_url, attachment, stored, thumbnail) 

520 

521 if isinstance(attachment.path, Path) and is_temp_path(attachment.path): 

522 attachment.path.unlink() 

523 attachment.path = None 

524 

525 if not self.__is_component: 

526 with self.xmpp.store.session(expire_on_commit=False) as orm: 

527 orm.add(stored) 

528 orm.commit() 

529 

530 msgs = self.__send_url( 

531 msg, 

532 legacy_msg_id, 

533 uploaded_url=new_url, 

534 caption=attachment.caption, 

535 carbon=carbon, 

536 when=when, 

537 correction=correction, 

538 **send_kwargs, 

539 ) 

540 if not self.__is_component and store_multi and legacy_msg_id: 

541 self.__store_multi(legacy_msg_id, msgs) 

542 

543 return new_url, msgs 

544 

545 def __send_body( 

546 self, 

547 body: str | None = None, 

548 legacy_msg_id: str | None = None, 

549 reply_to: MessageReference | None = None, 

550 when: datetime | None = None, 

551 thread: str | None = None, 

552 **kwargs: Any, # noqa:ANN401 

553 ) -> Message | None: 

554 if body: 

555 return self.send_text( 

556 body, 

557 legacy_msg_id, 

558 reply_to=reply_to, 

559 when=when, 

560 thread=thread, 

561 **kwargs, 

562 ) 

563 else: 

564 return None 

565 

566 async def send_files( 

567 self, 

568 attachments: Collection[LegacyAttachment], 

569 legacy_msg_id: str | None = None, 

570 body: str | None = None, 

571 *, 

572 reply_to: MessageReference | None = None, 

573 when: datetime | None = None, 

574 thread: str | None = None, 

575 body_first: bool = False, 

576 correction: bool = False, 

577 correction_event_id: str | None = None, 

578 **kwargs: Any, # noqa:ANN401 

579 ) -> None: 

580 # TODO: once the epic XEP-0385 vs XEP-0447 battle is over, pick 

581 # one and stop sending several attachments this way 

582 # we attach the legacy_message ID to the last message we send, because 

583 # we don't want several messages with the same ID (especially for MUC MAM) 

584 if not attachments and not body: 

585 # ignoring empty message 

586 return 

587 body_msg_id = ( 

588 legacy_msg_id if body_needs_msg_id(attachments, body, body_first) else None 

589 ) 

590 send_body = functools.partial( 

591 self.__send_body, 

592 body=body, 

593 reply_to=reply_to, 

594 when=when, 

595 thread=thread, 

596 correction=correction, 

597 legacy_msg_id=body_msg_id, 

598 correction_event_id=correction_event_id, 

599 **kwargs, 

600 ) 

601 all_msgs = [] 

602 if body_first: 

603 all_msgs.append(send_body()) 

604 for i, attachment in enumerate(attachments): 

605 legacy = legacy_msg_id if i == 0 and body_msg_id is None else None 

606 _url, msgs = await self.send_file( 

607 attachment, 

608 legacy, 

609 reply_to=reply_to, 

610 when=when, 

611 thread=thread, 

612 store_multi=False, 

613 **kwargs, 

614 ) 

615 all_msgs.extend(msgs) 

616 if not body_first: 

617 all_msgs.append(send_body()) 

618 self.__store_multi(legacy_msg_id, all_msgs) 

619 

620 def __store_multi( 

621 self, 

622 legacy_msg_id: str | None, 

623 all_msgs: Sequence[Message | None], 

624 ) -> None: 

625 if legacy_msg_id is None: 

626 return 

627 ids = [] 

628 for msg in all_msgs: 

629 if not msg: 

630 continue 

631 if stanza_id := msg.get_plugin("stanza_id", check=True): 

632 ids.append(stanza_id["id"]) 

633 else: 

634 ids.append(msg.get_id()) 

635 with self.xmpp.store.session() as orm: 

636 self.xmpp.store.id_map.set_msg( 

637 orm, self._recipient_pk(), str(legacy_msg_id), ids, self.is_participant 

638 ) 

639 orm.commit() 

640 

641 async def __get_thumbnail(self, attachment: LegacyAttachment) -> Thumbnail | None: 

642 if attachment.content_type is None: 

643 return None 

644 if not (attachment.data or attachment.stream or attachment.path): 

645 return None 

646 if not attachment.content_type.startswith("image"): 

647 return None 

648 

649 try: 

650 h, x, y = await self.xmpp.loop.run_in_executor( 

651 avatar_cache._thread_pool, 

652 get_thumbhash, 

653 attachment, 

654 ) 

655 except Exception as e: 

656 log.debug("Could not generate a thumbhash", exc_info=e) 

657 return None 

658 

659 thumbnail = Thumbnail() 

660 thumbnail["width"] = x 

661 thumbnail["height"] = y 

662 thumbnail["media-type"] = "image/thumbhash" 

663 thumbnail["uri"] = "data:image/thumbhash;base64," + urlquote(h) 

664 return thumbnail 

665 

666 

667class _AttachmentWithName(LegacyAttachment): 

668 name: str 

669 path: Path | None 

670 

671 

672class _AttachmentWithMetadata(_AttachmentWithName): 

673 size: int 

674 content_type: str 

675 

676 

677def _ensure_name(att: LegacyAttachment) -> _AttachmentWithName: 

678 if not att.name: 

679 if att.path: 

680 att.name = Path(att.path).name 

681 else: 

682 att.name = "unnamed-file" 

683 return cast(_AttachmentWithName, att) 

684 

685 

686async def _ensure_metadata(att: _AttachmentWithName) -> _AttachmentWithMetadata: 

687 if att.size is None: 

688 if att.data: 

689 att.size = len(att.data) 

690 elif att.stream: 

691 att.stream.seek(0, io.SEEK_END) 

692 att.size = att.stream.tell() 

693 att.stream.seek(0) 

694 elif att.path: 

695 assert isinstance(att.path, Path) 

696 att.size = att.path.stat().st_size 

697 elif att.url: 

698 async with ( 

699 aiohttp.ClientSession() as http, 

700 http.head(att.url) as resp, 

701 ): 

702 att.size = resp.content_length 

703 elif att.aio_stream: 

704 warnings.warn("A size should be passed with async iterators") 

705 tmp_dir = Path(tempfile.mkdtemp(prefix=_TEMP_PREFIX)) 

706 with (tmp_dir / att.name).open("wb") as fp: 

707 async for chunk in att.aio_stream: 

708 fp.write(chunk) 

709 att.path = Path(fp.name) 

710 att.size = att.path.stat().st_size 

711 

712 if not att.content_type: 

713 att.content_type, _encoding = guess_type(att.name) 

714 if not att.content_type: 

715 att.content_type = "application/octet-stream" 

716 

717 return cast(_AttachmentWithMetadata, att) 

718 

719 

720def body_needs_msg_id( 

721 attachments: Collection[LegacyAttachment], body: str | None, body_first: bool 

722) -> bool: 

723 if attachments: 

724 return bool(body and body_first) 

725 else: 

726 return True 

727 

728 

729def get_thumbhash(att: LegacyAttachment) -> tuple[str, int, int]: 

730 img, width, height = get_image(att) 

731 rgba_2d = list(img.get_flattened_data()) 

732 rgba = list(chain(*rgba_2d)) # type:ignore[arg-type,var-annotated] 

733 ints = thumbhash.rgba_to_thumb_hash(img.width, img.height, rgba) 

734 return base64.b64encode(bytes(ints)).decode(), width, height 

735 

736 

737def get_image(att: LegacyAttachment) -> tuple[Image.Image, int, int]: 

738 if att.data: 

739 img = Image.open(io.BytesIO(att.data)) 

740 elif att.path: 

741 with cast(Path, att.path).open("rb") as fp: 

742 img = Image.open(fp) 

743 img.load() 

744 elif att.stream: 

745 att.stream.seek(0) 

746 img = Image.open(att.stream) 

747 else: 

748 raise RuntimeError("No way to read the image") 

749 width, height = img.size 

750 rgba = img.convert("RGBA") 

751 if width > 100 or height > 100: 

752 rgba.thumbnail((100, 100)) 

753 return ImageOps.exif_transpose(rgba), width, height 

754 

755 

756def is_temp_path(path: Path) -> bool: 

757 try: 

758 return path.relative_to(_TEMP_ROOT).parts[0].startswith(_TEMP_PREFIX) 

759 except ValueError: 

760 return False 

761 

762 

763@contextlib.asynccontextmanager 

764async def _get_data( 

765 att: LegacyAttachment, http: aiohttp.ClientSession 

766) -> AsyncIterator[bytes | IO[bytes] | AsyncIterator[bytes]]: 

767 if (data := att.data or att.stream or att.aio_stream) is not None: 

768 yield data 

769 elif att.path is not None: 

770 assert isinstance(att.path, Path) 

771 with att.path.open("rb") as fp: 

772 yield fp 

773 elif att.url is not None: 

774 async with http.get(att.url) as resp_get: 

775 resp_get.raise_for_status() 

776 yield resp_get.content.iter_any() 

777 else: 

778 raise RuntimeError("NEVER") 

779 

780 

781_TEMP_ROOT = Path(tempfile.gettempdir()) 

782_TEMP_PREFIX = "slidge-async-iterator-download" 

783 

784 

785log = logging.getLogger(__name__)