Update main.py

main.py

@@ -6,6 +6,7 @@ import httpx
 import random
 import asyncio
 import urllib.parse
+import mimetypes
 from bs4 import BeautifulSoup
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import StreamingResponse
@@ -23,9 +24,8 @@ ACCESS_KEY = os.getenv("ACCESS_KEY", "0000")
 SERVICE_ACCOUNT_JSON_STR = os.getenv("GOOGLE_SERVICE_ACCOUNT_JSON")
 ua = UserAgent(fallback='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36')
 
-# Global Async Client
 client = httpx.AsyncClient(
-    timeout=httpx.Timeout(30.0, read=None), 
+    timeout=httpx.Timeout(60.0, read=None), 
     follow_redirects=True,
     limits=httpx.Limits(max_connections=100, max_keepalive_connections=20)
 )
@@ -55,106 +55,132 @@ 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
+
 @app.get("/download")
 async def download_proxy(request: Request, url: str, key: str = None):
     if key != ACCESS_KEY:
         raise HTTPException(status_code=403, detail="Access Denied")
 
-    # Double encoded ဖြစ်နေသော URL များကို ရှင်းလင်းရန်
-    url = urllib.parse.unquote(url)
+    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
 
-    range_header = request.headers.get('range')
+    # 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
 
-    # --- MediaFire & Others ---
-    if "drive.google.com" not in url:
-        target_link = None
-        headers = {'User-Agent': ua.random, 'Accept-Language': 'en-US,en;q=0.9'}
+    range_header = request.headers.get('range')
 
-        if "mediafire.com" in url:
-            try:
-                headers['Referer'] = 'https://www.mediafire.com/'
-                r = await client.get(url, headers=headers)
-                
-                # Regex ဖြင့် Direct Link ကို ရှာဖွေခြင်း
-                match = re.search(r'href="(https?://[a-zA-Z0-9-]+\.mediafire\.com/download/[^"]+)"', r.text)
-                if match:
-                    target_link = match.group(1)
-                else:
-                    soup = BeautifulSoup(r.text, 'html.parser')
-                    btn = soup.find('a', {'id': 'downloadButton'})
-                    if btn and btn.get('href'):
-                        # ပေးလာတဲ့ href မှာ https:// မပါရင် မူလ URL နဲ့ ပေါင်းပြီး အပြည့်အစုံဖြစ်အောင် ပြုလုပ်ခြင်း
-                        target_link = urllib.parse.urljoin(url, btn.get('href'))
-            except Exception as e:
-                print(f"MediaFire Error: {e}")
-        
-        elif "dropbox.com" in url:
-            target_link = url.replace("?dl=0", "").split("?")[0] + "?dl=1"
+    # --- MediaFire ---
+    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("'", "")
             
-        elif "1drv.ms" in url or "onedrive.live.com" in url:
-            encoded_url = base64.b64encode(bytes(url, 'utf-8')).decode('utf-8').replace('=', '').replace('/', '_').replace('+', '-')
-            target_link = f"https://api.onedrive.com/v1.0/shares/u!{encoded_url}/root/content"
-
-        if not target_link:
-            raise HTTPException(status_code=400, detail="Could not resolve direct link")
-        
-        return await stream_file(target_link, range_header)
+            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)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
 
     # --- Google Drive ---
-    file_id = get_google_file_id(url)
-    if not file_id:
-        raise HTTPException(status_code=400, detail="Invalid Google Drive Link")
-
-    accounts = get_all_service_accounts()
-    random.shuffle(accounts)
-
-    for account in accounts:
-        token = await get_token_for_account(account)
-        if not token: continue
-
-        api_link = f"https://www.googleapis.com/drive/v3/files/{file_id}?alt=media"
-        headers = {"Authorization": f"Bearer {token}"}
-        if range_header: headers['Range'] = range_header
+    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")
+
+        accounts = get_all_service_accounts()
+        random.shuffle(accounts)
+
+        for account in accounts:
+            token = await get_token_for_account(account)
+            if not token: continue
+            api_link = f"https://www.googleapis.com/drive/v3/files/{file_id}?alt=media"
+            headers = {"Authorization": f"Bearer {token}"}
+            if range_header: headers['Range'] = range_header
+            try:
+                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
 
-        try:
-            req = client.build_request("GET", api_link, headers=headers)
-            r = await client.send(req, stream=True)
-            if r.status_code in [200, 206]:
-                return await process_response(r)
-            await r.aclose()
-        except: continue
+        public_url = f"https://drive.google.com/uc?export=download&id={file_id}"
+        return await stream_file(public_url, range_header, filename)
 
-    # Public Fallback
-    public_url = f"https://drive.google.com/uc?export=download&id={file_id}"
-    return await stream_file(public_url, range_header)
+    else:
+        return await stream_file(clean_url, range_header, filename)
 
-async def stream_file(target_url, range_header):
+async def stream_file(target_url, range_header, filename, referer=None):
     headers = {'User-Agent': ua.random}
     if range_header: headers['Range'] = range_header
+    if referer: headers['Referer'] = referer
     
     try:
         req = client.build_request("GET", target_url, headers=headers)
         r = await client.send(req, stream=True)
-        return await process_response(r)
+        
+        # အကယ်၍ HTML ပြန်လာပါက ဗီဒီယိုမဟုတ်ဟု ယူဆပြီး Error ပြရန်
+        content_type = r.headers.get('content-type', '').lower()
+        if 'text/html' in content_type:
+            await r.aclose()
+            raise HTTPException(status_code=415, detail="Invalid Media: Received HTML instead of Video")
+            
+        return await process_response(r, filename)
     except Exception as e:
-        print(f"Stream Error for {target_url}: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
-async def process_response(r):
-    excluded = ['content-encoding', 'content-length', 'transfer-encoding', 'connection', 'content-disposition']
-    response_headers = {n: v for n, v in r.headers.items() if n.lower() not in excluded}
+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}
     
+    # အရေးကြီးသော Header များကို အတင်းအကျပ်သတ်မှတ်ခြင်း
+    response_headers['Content-Type'] = mime_type
     response_headers['Accept-Ranges'] = 'bytes'
-    if 'Content-Length' in r.headers:
-        response_headers['Content-Length'] = r.headers['Content-Length']
-    if 'Content-Range' in r.headers:
-        response_headers['Content-Range'] = r.headers['Content-Range']
-
-    response_headers['Content-Disposition'] = 'attachment'
+    
+    # nPlayer အတွက် filename header
+    quoted_filename = urllib.parse.quote(filename)
+    response_headers['Content-Disposition'] = f'inline; filename="{quoted_filename}"'
 
     async def stream_generator():
         try:
-            async for chunk in r.aiter_bytes(chunk_size=131072):
+            async for chunk in r.aiter_bytes(chunk_size=262144):
                 yield chunk
         finally:
             await r.aclose()
@@ -163,7 +189,7 @@ async def process_response(r):
         stream_generator(),
         status_code=r.status_code,
         headers=response_headers,
-        media_type='application/octet-stream'
+        media_type=mime_type
     )
 
 if __name__ == "__main__":