Update main.py

main.py

@@ -1,7 +1,6 @@
 import os
 import re
 import json
-import base64
 import httpx
 import random
 import asyncio
@@ -55,15 +54,13 @@ def get_google_file_id(url):
     fid = re.search(r'/(?:d|file/d|open\?id=)/([a-zA-Z0-9_-]+)', url)
     return fid.group(1) if fid else None
 
-def get_mime_from_filename(filename):
-    """ဖိုင်နာမည်ကိုကြည့်ပြီး Content-Type သတ်မှတ်ခြင်း"""
-    mime, _ = mimetypes.guess_type(filename)
-    if not mime:
-        if filename.lower().endswith('.mkv'): return 'video/x-matroska'
-        if filename.lower().endswith('.mp4'): return 'video/mp4'
-        if filename.lower().endswith('.avi'): return 'video/x-msvideo'
-        return 'video/mp4' # Default for players
-    return mime
+def get_clean_filename(url):
+    """URL ထဲကနေ ဖိုင်နာမည်ကို သန့်စင်စွာထုတ်ယူခြင်း"""
+    decoded_url = urllib.parse.unquote(url)
+    name = decoded_url.split('/')[-1].split('?')[0]
+    if not name or '.' not in name:
+        name = "video.mp4"
+    return name
 
 @app.get("/download")
 async def download_proxy(request: Request, url: str, key: str = None):
@@ -71,54 +68,52 @@ async def download_proxy(request: Request, url: str, key: str = None):
         raise HTTPException(status_code=403, detail="Access Denied")
 
     clean_url = urllib.parse.unquote(url)
-    while "%" in clean_url:
-        prev_url = clean_url
-        clean_url = urllib.parse.unquote(clean_url)
-        if prev_url == clean_url: break
-
-    # URL ထဲမှ ဖိုင်နာမည်ကို ထုတ်ယူခြင်း
-    filename = "video.mp4"
-    path_parts = clean_url.split('/')
-    for part in reversed(path_parts):
-        if '.' in part and len(part) > 4:
-            filename = part.split('?')[0]
-            break
-
+    filename = get_clean_filename(clean_url)
     range_header = request.headers.get('range')
 
-    # --- MediaFire ---
+    # --- MediaFire Section ---
     if "mediafire.com" in clean_url:
         try:
-            r = await client.get(clean_url, headers={'User-Agent': ua.random})
-            if r.status_code != 200:
-                raise HTTPException(status_code=r.status_code, detail="MediaFire page error")
-
-            target_link = None
-            match = re.search(r"https?://download[0-9]+\.mediafire\.com/[a-zA-Z0-9_-]+/[a-zA-Z0-9_-]+/[^\s'\"]+", r.text)
-            if match:
-                target_link = match.group(0).replace('"', '').replace("'", "")
-            
-            if not target_link:
-                soup = BeautifulSoup(r.text, 'html.parser')
-                btn = soup.find('a', {'id': 'downloadButton'})
-                if btn: target_link = 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}"
-            
-            if not target_link:
-                raise HTTPException(status_code=404, detail="Direct link extraction failed")
-            
-            return await stream_file(target_link, range_header, filename, referer=clean_url)
+            # Browser အသွင်ဖြင့် စာမျက်နှာကိုအရင်ဖတ်ခြင်း
+            async with httpx.AsyncClient(headers={'User-Agent': ua.random}, 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 page access failed")
+                
+                html_content = page_res.text
+                target_link = None
+
+                # Method A: Direct Search via Regex (The most reliable way)
+                # MediaFire direct links usually look like: https://download...
+                regex_match = re.search(r"https?://download[0-9]+\.mediafire\.com/[a-zA-Z0-9_-]+/[a-zA-Z0-9_-]+/[^\s'\"]+", html_content)
+                if regex_match:
+                    target_link = regex_match.group(0).replace('"', '').replace("'", "")
+                
+                # Method B: BeautifulSoup parsing
+                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 and download_btn.get('href'):
+                        target_link = download_btn.get('href')
+
+                if target_link:
+                    # Fix Relative URLs
+                    if target_link.startswith("//"): target_link = f"https:{target_link}"
+                    elif target_link.startswith("/"): target_link = f"https://www.mediafire.com{target_link}"
+                    
+                    return await stream_file(target_link, range_header, filename, referer=clean_url)
+                else:
+                    print(f"Extraction Failed for: {clean_url}")
+                    raise HTTPException(status_code=404, detail="Could not extract direct download link")
+
         except Exception as e:
+            print(f"MediaFire Logic Error: {str(e)}")
             raise HTTPException(status_code=500, detail=str(e))
 
-    # --- Google Drive ---
+    # --- Google Drive Section ---
     elif "drive.google.com" in clean_url:
         file_id = get_google_file_id(clean_url)
-        if not file_id:
-            raise HTTPException(status_code=400, detail="Invalid Google Drive Link")
+        if not file_id: raise HTTPException(status_code=400, detail="Invalid GD Link")
 
         accounts = get_all_service_accounts()
         random.shuffle(accounts)
@@ -130,14 +125,15 @@ async def download_proxy(request: Request, url: str, key: str = None):
             headers = {"Authorization": f"Bearer {token}"}
             if range_header: headers['Range'] = range_header
             try:
+                # build_request ကိုသုံးပြီး stream လုပ်ခြင်း
                 req = client.build_request("GET", api_link, headers=headers)
                 r = await client.send(req, stream=True)
                 if r.status_code in [200, 206]:
-                    # Google API က Content-Type အမှန်ပြန်ပေးလေ့ရှိသည်
                     return await process_response(r, filename)
                 await r.aclose()
             except: continue
 
+        # Fallback to Public Link
         public_url = f"https://drive.google.com/uc?export=download&id={file_id}"
         return await stream_file(public_url, range_header, filename)
 
@@ -153,30 +149,35 @@ async def stream_file(target_url, range_header, filename, referer=None):
         req = client.build_request("GET", target_url, headers=headers)
         r = await client.send(req, stream=True)
         
-        # အကယ်၍ HTML ပြန်လာပါက ဗီဒီယိုမဟုတ်ဟု ယူဆပြီး Error ပြရန်
-        content_type = r.headers.get('content-type', '').lower()
-        if 'text/html' in content_type:
+        # HTML ပြန်လာရင် ဗီဒီယိုမဟုတ်ကြောင်း သေချာစေရန်
+        ctype = r.headers.get('content-type', '').lower()
+        if 'text/html' in ctype:
             await r.aclose()
-            raise HTTPException(status_code=415, detail="Invalid Media: Received HTML instead of Video")
+            raise HTTPException(status_code=415, detail="Invalid Media Type: HTML Received")
             
         return await process_response(r, filename)
     except Exception as e:
+        print(f"Stream Error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
 async def process_response(r, filename):
-    mime_type = get_mime_from_filename(filename)
-    
-    # Player အတွက် လိုအပ်သော Header များကို သန့်စင်ခြင်း
-    safe_headers = ['content-length', 'content-range', 'accept-ranges']
-    response_headers = {n: v for n, v in r.headers.items() if n.lower() in safe_headers}
+    # MIME Type ကို အတင်းအကျပ် သတ်မှတ်ခြင်း
+    mime_type, _ = mimetypes.guess_type(filename)
+    if not mime_type or 'application' in mime_type:
+        if filename.lower().endswith('.mkv'): mime_type = 'video/x-matroska'
+        elif filename.lower().endswith('.mp4'): mime_type = 'video/mp4'
+        else: mime_type = 'video/mp4' # Player တွေအတွက် Safe အဖြစ်ဆုံး type
+
+    # nPlayer/VLC အတွက် လိုအပ်သော Header သန့်စင်မှု
+    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}
     
-    # အရေးကြီးသော Header များကို အတင်းအကျပ်သတ်မှတ်ခြင်း
-    response_headers['Content-Type'] = mime_type
-    response_headers['Accept-Ranges'] = 'bytes'
+    res_headers['Content-Type'] = mime_type
+    res_headers['Accept-Ranges'] = 'bytes'
     
-    # nPlayer အတွက် filename header
-    quoted_filename = urllib.parse.quote(filename)
-    response_headers['Content-Disposition'] = f'inline; filename="{quoted_filename}"'
+    # Filename Header
+    quoted_name = urllib.parse.quote(filename)
+    res_headers['Content-Disposition'] = f'inline; filename="{quoted_name}"'
 
     async def stream_generator():
         try:
@@ -188,7 +189,7 @@ async def process_response(r, filename):
     return StreamingResponse(
         stream_generator(),
         status_code=r.status_code,
-        headers=response_headers,
+        headers=res_headers,
         media_type=mime_type
     )