From 9ebc7f5d57875b1dd85b917ac9a58cca3974c6ca Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sun, 28 Jan 2018 21:35:04 +0200
Subject: [PATCH] Only send PNGs as images

---
 mautrix_telegram/portal.py | 19 +++++++++++++++----
 mautrix_telegram/user.py   |  2 +-
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py
index 8e21a198..53b14dc0 100644
--- a/mautrix_telegram/portal.py
+++ b/mautrix_telegram/portal.py
@@ -26,6 +26,8 @@ import magic
 from .db import Portal as DBPortal, Message as DBMessage
 from . import puppet as p, user as u, formatter
 
+mimetypes.init()
+
 config = None
 
 
@@ -233,6 +235,18 @@ class Portal:
     # endregion
     # region Matrix event handling
 
+    def _get_file_meta(self, body, mime):
+        file_name = None
+        try:
+            current_extension = body[body.rindex("."):]
+            if mimetypes.types_map[current_extension] == mime:
+                file_name = body
+            else:
+                file_name = f"matrix_upload{mimetypes.guess_extension(mime)}"
+        except (ValueError, KeyError):
+            file_name = f"matrix_upload{mimetypes.guess_extension(mime)}"
+        return file_name, None if file_name == body else body
+
     def handle_matrix_message(self, sender, message, event_id):
         type = message["msgtype"]
         if type == "m.text":
@@ -252,12 +266,9 @@ class Portal:
             file = self.main_intent.download_file(message["url"])
 
             info = message["info"]
-            body = message["body"]
             mime = info["mimetype"]
 
-            extension = mimetypes.guess_extension(mime)
-            file_name = body if body.endswith(extension) else f"matrix_upload{extension}"
-            caption = None if file_name == body else body
+            file_name, caption = self._get_file_meta(message["body"], mime)
 
             attributes = [DocumentAttributeFilename(file_name=file_name)]
             if "w" in info and "h" in info:
diff --git a/mautrix_telegram/user.py b/mautrix_telegram/user.py
index c771ecd1..dc128134 100644
--- a/mautrix_telegram/user.py
+++ b/mautrix_telegram/user.py
@@ -147,7 +147,7 @@ class User:
 
         file_handle = self.client.upload_file(file, file_name=file_name, use_cache=False)
 
-        if mime_type.startswith("image/"):
+        if mime_type == "image/png":
             media = InputMediaUploadedPhoto(file_handle, caption or "")
         else:
             attr_dict = {type(attr): attr for attr in attributes}