From 33b9189b90bb2456ece23ab5149428ce071f28f6 Mon Sep 17 00:00:00 2001
From: Hirusha Himath <git.itzfork@gmail.com>
Date: Sat, 15 Jun 2024 10:13:03 +0530
Subject: [PATCH] Possible fix for #64 (#65)

improvements on file name finder on lib/ddl
---
 .gitignore        |  4 +++-
 megadl/lib/ddl.py | 23 ++++++++++++++++++++++-
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index 0027a9ad..d61c5a9b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,8 @@
 # dev environment
 .env
 venv
+.venv
 .mega/
 *.session
-NexaBots
\ No newline at end of file
+NexaBots
+tests
\ No newline at end of file
diff --git a/megadl/lib/ddl.py b/megadl/lib/ddl.py
index 2b0151b8..ae0bd08d 100644
--- a/megadl/lib/ddl.py
+++ b/megadl/lib/ddl.py
@@ -6,6 +6,8 @@
 import os
 import re
 import asyncio
+import hashlib
+import mimetypes
 
 from time import time
 from aiohttp import ClientSession
@@ -19,6 +21,7 @@
     r"https://drive\.google\.com/file/d/(.*?)/.*?\?usp=(sharing|drive_link)"
 )
 STR_GD_REPLACE = r"https://drive.google.com/uc?export=download&id=\1"
+DEFAULT_EXT = ".bin"
 
 
 class Downloader:
@@ -81,7 +84,6 @@ async def from_ddl(
         # Create folder if it doesn't exist
         wpath = f"{path}/{chat_id}"
         os.makedirs(wpath)
-        wpath = f"{wpath}/{os.path.basename(url)}"
 
         async with ClientSession() as session:
             _chunksize = int(os.getenv("CHUNK_SIZE"))
@@ -89,6 +91,25 @@ async def from_ddl(
                 # Raise HttpStatusError on failed requests
                 if resp.status != 200:
                     raise HttpStatusError(resp.status)
+
+                # Try to get the file name
+                cnt_disp = resp.headers.get("Content-Disposition")
+                if (
+                    cnt_disp
+                    and (cd_parts := cnt_disp.split("filename="))
+                    and len(cd_parts) > 1
+                ):
+                    fname = cd_parts[1].strip('\"')
+
+                elif _ftype := mimetypes.guess_type(url):
+                    _fext = mimetypes.guess_extension(_ftype)
+                    fname = f"{hashlib.md5(url.encode()).hexdigest()}{DEFAULT_EXT if not _fext else _fext}"
+
+                else:
+                    fname = os.path.basename(url)
+
+                wpath = f"{wpath}/{fname}"
+
                 # Handle content length header
                 total = resp.content_length
                 curr = 0