From c6fd3d999f88d7f7ee5d351f39b4573297374e8a Mon Sep 17 00:00:00 2001
From: Dawid Rejowski <dawid@rejowski.xyz>
Date: Sun, 11 Sep 2022 23:12:24 +0200
Subject: [PATCH] Automatic detecting mime type of images

---
 requirements.txt                   | 1 +
 stickerbridge/chat_functions.py    | 4 +++-
 stickerbridge/matrix_reuploader.py | 4 ++--
 stickerbridge/sticker_types.py     | 3 +--
 stickerbridge/telegram_exporter.py | 3 +--
 5 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index cad4a2f..cc607c4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,5 +3,6 @@ pillow
 matrix-nio
 pyyaml
 aiofiles
+python-magic
 lottie
 cairosvg
\ No newline at end of file
diff --git a/stickerbridge/chat_functions.py b/stickerbridge/chat_functions.py
index e335e48..7542b75 100644
--- a/stickerbridge/chat_functions.py
+++ b/stickerbridge/chat_functions.py
@@ -1,6 +1,7 @@
 import os
 
 import aiofiles.os
+import magic
 import logging
 
 from nio import AsyncClient, UploadResponse, ErrorResponse, RoomGetStateEventError
@@ -52,7 +53,8 @@ async def upload_stickerpack(client: AsyncClient, room_id: str, stickerset: Matr
     return await client.room_put_state(room_id, 'im.ponies.room_emotes', stickerset.json(), state_key=stickerset.name())
 
 
-async def upload_image(client: AsyncClient, image: str, mime_type: str):
+async def upload_image(client: AsyncClient, image: str):
+    mime_type = magic.from_file(image, mime=True)
     file_stat = await aiofiles.os.stat(image)
     async with aiofiles.open(image, "r+b") as f:
         resp, maybe_keys = await client.upload(
diff --git a/stickerbridge/matrix_reuploader.py b/stickerbridge/matrix_reuploader.py
index ba93fea..33cc667 100644
--- a/stickerbridge/matrix_reuploader.py
+++ b/stickerbridge/matrix_reuploader.py
@@ -46,9 +46,9 @@ class MatrixReuploader:
         converted_stickerset = await self.exporter.get_stickerset(stickerset.name())
         yield self.STATUS_UPLOADING
         for sticker in converted_stickerset:
-            with tempfile.NamedTemporaryFile('w+b', suffix=sticker.mime_type.replace('image/', '')) as file:
+            with tempfile.NamedTemporaryFile('w+b') as file:
                 file.write(sticker.image_data)
-                sticker_mxc = await upload_image(self.client, file.name, sticker.mime_type)
+                sticker_mxc = await upload_image(self.client, file.name)
             stickerset.add_sticker(sticker_mxc, sticker.alt_text)
 
         if not stickerset.count():
diff --git a/stickerbridge/sticker_types.py b/stickerbridge/sticker_types.py
index b35e57a..367e92c 100644
--- a/stickerbridge/sticker_types.py
+++ b/stickerbridge/sticker_types.py
@@ -1,9 +1,8 @@
 class Sticker:
     """Custom type for easier transfering sticker data between functions and classes with simple lists and returns"""
-    def __init__(self, image_data, alt_text: str, mime_type: str):
+    def __init__(self, image_data, alt_text: str):
         self.image_data = image_data
         self.alt_text = alt_text
-        self.mime_type = mime_type
 
 
 class MatrixStickerset:
diff --git a/stickerbridge/telegram_exporter.py b/stickerbridge/telegram_exporter.py
index 08eaaa8..99f0948 100644
--- a/stickerbridge/telegram_exporter.py
+++ b/stickerbridge/telegram_exporter.py
@@ -71,9 +71,8 @@ class TelegramExporter:
             raw_data = await self.client.download_media(sticker_document, file=bytes)
             if sticker_document.mime_type == 'image/webp':
                 data, width, height = _convert_image(raw_data)
-                result.append(Sticker(data, alt, 'image/png'))
             if sticker_document.mime_type == 'application/x-tgsticker':
                 data = _convert_animation(raw_data)
-                result.append(Sticker(data, alt, 'image/webp'))
+            result.append(Sticker(data, alt))
 
         return result