Make mime type extension guessing saner

mautrix_telegram/portal.py

@@ -69,7 +69,7 @@ from mautrix_appservice import MatrixRequestError, IntentError, AppService, Inte
 from .types import MatrixEventID, MatrixRoomID, MatrixUserID, TelegramID
 from .context import Context
 from .db import Portal as DBPortal, Message as DBMessage, TelegramFile as DBTelegramFile
-from .util import ignore_coro
+from .util import ignore_coro, sane_mimetypes
 from . import puppet as p, user as u, formatter, util
 
 if TYPE_CHECKING:
@@ -78,8 +78,6 @@ if TYPE_CHECKING:
     from .config import Config
     from .tgclient import MautrixTelegramClient
 
-mimetypes.init()
-
 config = None  # type: Config
 
 TypeMessage = Union[Message, MessageService]
@@ -783,14 +781,9 @@ class Portal:
                 return body + current_extension
         except (ValueError, KeyError):
             pass
-        ext_override = {
-            "image/jpeg": ".jpg"
-        }
         if mime:
-            ext = ext_override.get(mime, mimetypes.guess_extension(mime))
-            return f"matrix_upload{ext}"
-        else:
-            return ""
+            return f"matrix_upload{sane_mimetypes.guess_extension(mime)}"
+        return ""
 
     def get_config(self, key: str) -> Any:
         local = util.recursive_get(self.local_config, key)
@@ -1228,7 +1221,7 @@ class Portal:
 
         file = await self.main_intent.download_file(url)
         mime = magic.from_buffer(file, mime=True)
-        ext = mimetypes.guess_extension(mime)
+        ext = sane_mimetypes.guess_extension(mime)
         uploaded = await sender.client.upload_file(file, file_name=f"avatar{ext}", use_cache=False)
         photo = InputChatUploadedPhoto(file=uploaded)
 
@@ -1421,11 +1414,7 @@ class Portal:
             "orientation": 0,
             "mimetype": file.mime_type,
         }
-        ext_override = {
-            "image/jpeg": ".jpg"
-        }
-        name = "image" + ext_override.get(file.mime_type,
-                                          mimetypes.guess_extension(file.mime_type))
+        name = f"image{sane_mimetypes.guess_extension(file.mime_type)}"
         await intent.set_typing(self.mxid, is_typing=False)
         result = await intent.send_image(self.mxid, file.mxc, info=info, text=name,
                                          relates_to=relates_to, timestamp=evt.date,

mautrix_telegram/util/__init__.py

@@ -3,6 +3,5 @@ from .format_duration import format_duration
 from .signed_token import sign_token, verify_token
 from .recursive_dict import recursive_del, recursive_set, recursive_get
 
-
 def ignore_coro(coro):
     pass

mautrix_telegram/util/file_transfer.py

@@ -32,6 +32,7 @@ from mautrix_appservice import IntentAPI
 
 from ..tgclient import MautrixTelegramClient
 from ..db import TelegramFile as DBTelegramFile
+from ..util import sane_mimetypes
 
 try:
     from PIL import Image
@@ -121,7 +122,7 @@ async def transfer_thumbnail_to_matrix(client: MautrixTelegramClient, intent: In
     if db_file:
         return db_file
 
-    video_ext = mimetypes.guess_extension(mime)
+    video_ext = sane_mimetypes.guess_extension(mime)
     if VideoFileClip and video_ext:
         try:
             file, width, height = _read_video_thumbnail(video, video_ext, frame_ext="png")

mautrix_telegram/util/sane_mimetypes.py

@@ -0,0 +1,38 @@
+# -*- coding: future_fstrings -*-
+# mautrix-telegram - A Matrix-Telegram puppeting bridge
+# Copyright (C) 2019 Tulir Asokan
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+import mimetypes
+
+mimetypes.init()
+
+sanity_overrides = {
+    "image/jpeg": ".jpeg",
+    "image/tiff": ".tiff",
+    "text/plain": ".txt",
+    "text/html": ".html",
+    "audio/mpeg": ".mp3",
+    "audio/ogg": ".ogg",
+    "application/xml": ".xml",
+    "application/octet-stream": "",
+    "application/x-msdos-program": ".exe",
+}
+
+
+def guess_extension(mime: str) -> str:
+    try:
+        return sanity_overrides[mime]
+    except KeyError:
+        return mimetypes.guess_extension(mime)