Fixed download with weird file name

legzdev · legzdev · commit bee328920ca6 · 2023-01-10T19:50:08.000-05:00
diff --git a/WebStreamer/server/stream_routes.py b/WebStreamer/server/stream_routes.py
@@ -12,6 +12,7 @@
 from WebStreamer.bot import multi_clients, work_loads
 from WebStreamer.server.exceptions import FIleNotFound, InvalidHash
 from WebStreamer import Var, utils, StartTime, __version__, StreamBot
+from WebStreamer.utils import get_name
 
 logger = logging.getLogger("routes")
 
@@ -116,21 +117,11 @@ async def media_streamer(request: web.Request, message_id: int, secure_hash: str
         file_id, index, offset, first_part_cut, last_part_cut, part_count, chunk_size
     )
     mime_type = file_id.mime_type
-    file_name = file_id.file_name
+    file_name = get_name(file_id)
     disposition = "attachment"
 
-    if mime_type:
-        if not file_name:
-            try:
-                file_name = f"{secrets.token_hex(2)}.{mime_type.split('/')[1]}"
-            except (IndexError, AttributeError):
-                file_name = f"{secrets.token_hex(2)}.unknown"
-    else:
-        if file_name:
-            mime_type = mimetypes.guess_type(file_id.file_name)
-        else:
-            mime_type = "application/octet-stream"
-            file_name = f"{secrets.token_hex(2)}.unknown"
+    if not mime_type:
+        mime_type = mimetypes.guess_type(file_name)[0] or "application/octet-stream"
 
     if "video/" in mime_type or "audio/" in mime_type or "/html" in mime_type:
         disposition = "inline"
diff --git a/WebStreamer/utils/file_properties.py b/WebStreamer/utils/file_properties.py
@@ -58,12 +58,17 @@ def get_hash(media_msg: Union[str, Message], length: int) -> str:
     return long_hash[:length]
 
 
-def get_name(media_msg: Message) -> str:
-    media = get_media_from_message(media_msg)
-    file_name = getattr(media, 'file_name', "") or ""
+def get_name(media_msg: Message | FileId) -> str:
+
+    if isinstance(media_msg, Message):
+        media = get_media_from_message(media_msg)
+        file_name = getattr(media, "file_name", "")
+
+    elif isinstance(media_msg, FileId):
+        file_name = getattr(media_msg, "file_name", "")
 
     if not file_name:
-        if media_msg.media and media_msg.media.value:
+        if isinstance(media_msg, Message) and media_msg.media:
             media_type = media_msg.media.value
         else:
             media_type = "file"
