From be477065788b14527182127fbc8fc3f6eb0028e7 Mon Sep 17 00:00:00 2001
From: Nick Mills-Barrett <nick@beeper.com>
Date: Mon, 25 Mar 2024 14:31:28 +0000
Subject: [PATCH] Implement bot/channel file size limit

---
 .../portal_util/message_convert.py            | 104 +++++++++++++++++-
 mautrix_telegram/util/__init__.py             |   1 +
 2 files changed, 103 insertions(+), 2 deletions(-)

diff --git a/mautrix_telegram/portal_util/message_convert.py b/mautrix_telegram/portal_util/message_convert.py
index 8b3218f7..dbff006a 100644
--- a/mautrix_telegram/portal_util/message_convert.py
+++ b/mautrix_telegram/portal_util/message_convert.py
@@ -170,8 +170,13 @@ async def convert(
         if not client:
             client = source.client
         if hasattr(evt, "media") and isinstance(evt.media, self._allowed_media):
-            convert_media = self._media_converters[type(evt.media)]
-            converted = await convert_media(source=source, intent=intent, evt=evt, client=client)
+            if self._should_convert_full_document(evt.media, is_bot, is_channel):
+                convert_media = self._media_converters[type(evt.media)]
+                converted = await convert_media(
+                    source=source, intent=intent, evt=evt, client=client
+                )
+            else:
+                converted = await self._convert_document_thumb_only(source, intent, evt, client)
         elif evt.message:
             converted = await self._convert_text(source, intent, is_bot, evt, client)
         else:
@@ -204,6 +209,16 @@ async def convert(
             )
         return converted
 
+    def _should_convert_full_document(self, media, is_bot: bool, is_channel: bool) -> bool:
+        if not isinstance(media, MessageMediaDocument):
+            return True
+        size = media.document.size
+        if is_bot and self.config["bridge.document_as_link_size.bot"]:
+            return size < self.config["bridge.document_as_link_size.bot"] * 1000**2
+        if is_channel and self.config["bridge.document_as_link_size.channel"]:
+            return size < self.config["bridge.document_as_link_size.channel"] * 1000**2
+        return True
+
     @staticmethod
     def _caption_to_message(converted: ConvertedMessage) -> None:
         content, caption = converted.content, converted.caption
@@ -493,6 +508,91 @@ def _adjust_ttl(ttl: int | None) -> int | None:
             # but we can only count it from read receipt.
             return ttl * 5
 
+    async def _convert_document_thumb_only(
+        self,
+        source: au.AbstractUser,
+        intent: IntentAPI,
+        evt: Message,
+        client: MautrixTelegramClient,
+    ) -> ConvertedMessage | None:
+        document = evt.media.document
+
+        if not document:
+            return None
+
+        external_link_content = "Unsupported file, please access directly on Telegram"
+
+        external_url = self._get_external_url(evt)
+        # We don't generate external URLs for bot users so only set if known
+        if external_url is not None:
+            external_link_content = (
+                f"Unsupported file, please access directly on Telegram here: {external_url}"
+            )
+
+        attrs = _parse_document_attributes(document.attributes)
+        file = None
+
+        thumb_loc, thumb_size = self.get_largest_photo_size(document)
+        if thumb_size and not isinstance(thumb_size, (PhotoSize, PhotoCachedSize)):
+            self.log.debug(f"Unsupported thumbnail type {type(thumb_size)}")
+            thumb_loc = None
+            thumb_size = None
+        if thumb_loc:
+            try:
+                file = await util.transfer_thumbnail_to_matrix(
+                    client,
+                    intent,
+                    thumb_loc,
+                    video=None,
+                    mime_type=document.mime_type,
+                    encrypt=self.portal.encrypted,
+                    async_upload=self.config["homeserver.async_media"],
+                )
+            except Exception:
+                self.log.exception("Failed to transfer thumbnail")
+        if not file:
+            name = attrs.name or ""
+            caption = f"\n{evt.message}" if evt.message else ""
+            return ConvertedMessage(
+                content=TextMessageEventContent(
+                    msgtype=MessageType.NOTICE,
+                    body=f"{name}{caption}\n{external_link_content}",
+                )
+            )
+
+        info, name = _parse_document_meta(evt, file, attrs, thumb_size)
+
+        event_type = EventType.ROOM_MESSAGE
+        if not name:
+            ext = sane_mimetypes.guess_extension(file.mime_type) or ""
+            name = "unnamed_file" + ext
+
+        content = MediaMessageEventContent(
+            body=name,
+            info=info,
+            msgtype={
+                "video/": MessageType.VIDEO,
+                "audio/": MessageType.AUDIO,
+                "image/": MessageType.IMAGE,
+            }.get(info.mimetype[:6], MessageType.FILE),
+        )
+        if file.decryption_info:
+            content.file = file.decryption_info
+        else:
+            content.url = file.mxc
+
+        caption_content = (
+            await formatter.telegram_to_matrix(evt, source, client) if evt.message else None
+        )
+        caption_content = f"{caption_content}\n{external_link_content}"
+
+        return ConvertedMessage(
+            type=event_type,
+            content=content,
+            caption=caption_content,
+            disappear_seconds=self._adjust_ttl(evt.media.ttl_seconds),
+        )
+
     async def _convert_document(
         self,
         source: au.AbstractUser,
diff --git a/mautrix_telegram/util/__init__.py b/mautrix_telegram/util/__init__.py
index eb5a7e21..53f6120e 100644
--- a/mautrix_telegram/util/__init__.py
+++ b/mautrix_telegram/util/__init__.py
@@ -4,6 +4,7 @@
     convert_image,
     transfer_custom_emojis_to_matrix,
     transfer_file_to_matrix,
+    transfer_thumbnail_to_matrix,
     unicode_custom_emoji_map,
 )
 from .parallel_file_transfer import parallel_transfer_to_telegram