Update main.py

main.py

@@ -60,10 +60,11 @@ def get_google_file_id(url):
     return fid.group(1) if fid else None
 
 def get_clean_filename(url):
+    """URL ထဲမှ ဖိုင်နာမည်နှင့် Extension (.mp4/.mkv) ကို တိကျစွာထုတ်ယူခြင်း"""
     decoded_url = urllib.parse.unquote(url)
     name = decoded_url.split('/')[-1].split('?')[0]
     if not name or '.' not in name:
-        name = "video.mp4"
+        return "video.mp4"
     return name
 
 @app.get("/download")
@@ -76,48 +77,46 @@ async def download_proxy(request: Request, url: str, key: str = None):
     range_header = request.headers.get('range')
     current_time = time.time()
 
-    # --- MediaFire Section (Modified to use Redirect) ---
+    # --- MediaFire Section ---
     if "mediafire.com" in clean_url:
         cached_data = MEDIAFIRE_CACHE.get(clean_url)
         
-        # Cache ရှိရင် တိုက်ရိုက် Redirect လုပ်မည်
+        target_link = None
         if cached_data and (current_time - cached_data['time']) < CACHE_TTL:
-            return RedirectResponse(url=cached_data['link'])
-
-        try:
-            req_ua = ua.random
-            async with httpx.AsyncClient(headers={'User-Agent': req_ua}, follow_redirects=True) as temp_client:
-                page_res = await temp_client.get(clean_url)
-                if page_res.status_code != 200:
-                    raise HTTPException(status_code=page_res.status_code, detail="MediaFire access failed")
-                
-                html_content = page_res.text
-                target_link = None
-
-                match = re.search(r"https?://download[0-9]+\.mediafire\.com/[a-zA-Z0-9_-]+/[a-zA-Z0-9_-]+/[^\s'\"]+", html_content)
-                if match:
-                    target_link = match.group(0).replace('"', '').replace("'", "")
-                
-                if not target_link:
-                    soup = BeautifulSoup(html_content, 'html.parser')
-                    download_btn = soup.find('a', {'id': 'downloadButton'}) or soup.find('a', {'class': 'input_btn_p'})
-                    if download_btn: target_link = download_btn.get('href')
-
-                if target_link:
-                    if target_link.startswith("//"): target_link = f"https:{target_link}"
-                    elif target_link.startswith("/"): target_link = f"https://www.mediafire.com{target_link}"
+            target_link = cached_data['link']
+        else:
+            try:
+                req_ua = ua.random
+                async with httpx.AsyncClient(headers={'User-Agent': req_ua}, follow_redirects=True) as temp_client:
+                    page_res = await temp_client.get(clean_url)
+                    if page_res.status_code != 200:
+                        raise HTTPException(status_code=page_res.status_code, detail="MediaFire access failed")
                     
-                    # သိမ်းဆည်းခြင်း
-                    MEDIAFIRE_CACHE[clean_url] = {'link': target_link, 'time': current_time}
+                    html_content = page_res.text
+                    # Direct link Regex
+                    match = re.search(r"https?://download[0-9]+\.mediafire\.com/[a-zA-Z0-9_-]+/[a-zA-Z0-9_-]+/[^\s'\"]+", html_content)
+                    if match:
+                        target_link = match.group(0).replace('"', '').replace("'", "")
                     
-                    # Proxy အစား RedirectResponse ကိုသုံးပြီး Player ဆီ လွှဲပေးလိုက်ခြင်း
-                    return RedirectResponse(url=target_link)
-                else:
-                    raise HTTPException(status_code=404, detail="Direct link extraction failed")
-
-        except Exception as e:
-            print(f"MediaFire Error: {e}")
-            raise HTTPException(status_code=500, detail=str(e))
+                    if not target_link:
+                        soup = BeautifulSoup(html_content, 'html.parser')
+                        download_btn = soup.find('a', {'id': 'downloadButton'}) or soup.find('a', {'class': 'input_btn_p'})
+                        if download_btn: target_link = download_btn.get('href')
+
+                    if target_link:
+                        if target_link.startswith("//"): target_link = f"https:{target_link}"
+                        elif target_link.startswith("/"): target_link = f"https://www.mediafire.com{target_link}"
+                        MEDIAFIRE_CACHE[clean_url] = {'link': target_link, 'time': current_time}
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+
+        if target_link:
+            # .mp4 သို့မဟုတ် .mkv ဖြစ်ကြောင်း Player သိစေရန် Extension hint ထည့်ခြင်း
+            separator = "&" if "?" in target_link else "?"
+            hinted_link = f"{target_link}{separator}type=video&file={urllib.parse.quote(filename)}"
+            return RedirectResponse(url=hinted_link)
+        else:
+            raise HTTPException(status_code=404, detail="Direct link extraction failed")
 
     # --- Google Drive Section (Stay as Proxy) ---
     elif "drive.google.com" in clean_url:
@@ -158,9 +157,15 @@ async def stream_file(target_url, range_header, filename):
         raise HTTPException(status_code=500, detail=str(e))
 
 async def process_response(r, filename):
+    # MIME Type ကို .mp4 သို့မဟုတ် .mkv အလိုက် အလိုအလျောက် သတ်မှတ်ခြင်း
     mime_type, _ = mimetypes.guess_type(filename)
     if not mime_type or 'application' in mime_type:
-        mime_type = 'video/x-matroska' if filename.endswith('.mkv') else 'video/mp4'
+        if filename.lower().endswith('.mp4'):
+            mime_type = 'video/mp4'
+        elif filename.lower().endswith('.mkv'):
+            mime_type = 'video/x-matroska'
+        else:
+            mime_type = 'video/mp4'
 
     safe_headers = ['content-length', 'content-range', 'accept-ranges', 'cache-control']
     res_headers = {n: v for n, v in r.headers.items() if n.lower() in safe_headers}