Update main.py

main.py

@@ -1,12 +1,10 @@
 import os
 import re
 import json
-import time
+import base64
 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
@@ -24,19 +22,16 @@ 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')
 
-# MediaFire အတွက် Cache Memory
-MEDIAFIRE_CACHE = {}
-CACHE_TTL = 1800 # Cache သက်တမ်း စက္ကန့် 1800 (မိနစ် 30)
-
+# Global Async Client: Video ကျော်ကြည့်တဲ့အခါ Connection အသစ်တွေ မြန်မြန်ဆောက်နိုင်ဖို့
 client = httpx.AsyncClient(
-    timeout=httpx.Timeout(60.0, read=None), 
+    timeout=httpx.Timeout(20.0, read=None), 
     follow_redirects=True,
     limits=httpx.Limits(max_connections=100, max_keepalive_connections=20)
 )
 
 @app.get("/")
 async def index():
-    return {"status": "Online"}
+    return {"Online"}
 
 def get_all_service_accounts():
     if not SERVICE_ACCOUNT_JSON_STR: return []
@@ -49,6 +44,7 @@ async def get_token_for_account(cred_dict):
     try:
         scopes = ['https://www.googleapis.com/auth/drive.readonly']
         creds = service_account.Credentials.from_service_account_info(cred_dict, scopes=scopes)
+        # Token refresh ကို thread ထဲမှာ run ပါတယ် (Google library က sync ဖြစ်လို့)
         loop = asyncio.get_event_loop()
         auth_req = google.auth.transport.requests.Request()
         await loop.run_in_executor(None, creds.refresh, auth_req)
@@ -59,162 +55,101 @@ 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_clean_filename(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):
     if key != ACCESS_KEY:
         raise HTTPException(status_code=403, detail="Access Denied")
 
-    clean_url = urllib.parse.unquote(url)
-    filename = get_clean_filename(clean_url)
     range_header = request.headers.get('range')
-    
-    req_ua = ua.random 
-    current_time = time.time()
 
-    # --- MediaFire Section ---
-    if "mediafire.com" in clean_url:
-        
-        # 1. Cache ထဲမှာရှိမရှိ အရင်စစ်ဆေးခြင်း (Player များ Seek လုပ်ရာတွင် လွယ်ကူစေရန်)
-        cached_data = MEDIAFIRE_CACHE.get(clean_url)
-        if cached_data and (current_time - cached_data['time']) < CACHE_TTL:
+    # --- 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'}
+
+        if "mediafire.com" in url:
             try:
-                return await stream_file(
-                    cached_data['link'], range_header, filename, 
-                    referer=clean_url, req_ua=cached_data['ua'], cookies=cached_data['cookies']
-                )
-            except HTTPException as e:
-                # အကယ်၍ Cached Link သည် သက်တမ်းကုန်သွား၍ 415 ပြန်လာပါက Cache ကိုဖျက်ပြီး အသစ်ပြန်ရှာမည်
-                if e.status_code == 415:
-                    print(f"Cache expired for {clean_url}, refetching...")
-                    del MEDIAFIRE_CACHE[clean_url]
-                else:
-                    raise
-                    
-        # 2. Cache မရှိပါက သို့မဟုတ် Expire ဖြစ်သွားပါက အသစ်ပြန်ရှာခြင်း
-        try:
-            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 page access failed")
-                
-                html_content = page_res.text
-                target_link = None
-                cookies = temp_client.cookies 
-
-                match = re.search(r"https?://download[0-9]+\.mediafire\.com/[a-zA-Z0-9_-]+/[a-zA-Z0-9_-]+/[^\s'\"]+", html_content)
+                headers['Referer'] = 'https://www.mediafire.com/'
+                r = await client.get(url, headers=headers)
+                # Regex ဖြင့် Direct Link ကို အမြန်ဆုံး ရှာဖွေခြင်း
+                match = re.search(r'href="(https?://download[^"]+)"', r.text)
                 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 and download_btn.get('href'):
-                        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}"
-                    
-                    # နောက်ထပ် Range Request များအတွက် Memory ထဲတွင် မှတ်သားထားခြင်း
-                    MEDIAFIRE_CACHE[clean_url] = {
-                        'link': target_link,
-                        'ua': req_ua,
-                        'cookies': dict(cookies),
-                        'time': current_time
-                    }
-                    
-                    return await stream_file(target_link, range_header, filename, referer=clean_url, req_ua=req_ua, cookies=cookies)
+                    target_link = match.group(1)
                 else:
-                    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 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 GD 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]:
-                    return await process_response(r, filename)
-                await r.aclose()
-            except: continue
+                    soup = BeautifulSoup(r.text, 'html.parser')
+                    btn = soup.find('a', {'id': 'downloadButton'})
+                    target_link = btn.get('href') if btn else None
+            except Exception as e:
+                print(f"MediaFire Error: {e}")
+        
+        elif "dropbox.com" in url:
+            target_link = url.replace("?dl=0", "").split("?")[0] + "?dl=1"
+            
+        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)
+
+    # --- Google Drive ---
+    file_id = get_google_file_id(url)
+    if not file_id:
+        raise HTTPException(status_code=400, detail="Invalid Google Drive Link")
 
-        public_url = f"https://drive.google.com/uc?export=download&id={file_id}"
-        return await stream_file(public_url, range_header, filename, req_ua=req_ua)
+    accounts = get_all_service_accounts()
+    random.shuffle(accounts)
 
-    else:
-        return await stream_file(clean_url, range_header, filename, req_ua=req_ua)
+    for account in accounts:
+        token = await get_token_for_account(account)
+        if not token: continue
 
-async def stream_file(target_url, range_header, filename, referer=None, req_ua=None, cookies=None):
-    headers = {'User-Agent': req_ua if req_ua else ua.random}
+        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:
+            # Stream mode ဖြင့် ချိတ်ဆက်ခြင်း
+            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 Fallback
+    public_url = f"https://drive.google.com/uc?export=download&id={file_id}"
+    return await stream_file(public_url, range_header)
+
+async def stream_file(target_url, range_header):
+    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, cookies=cookies)
+        req = client.build_request("GET", target_url, headers=headers)
         r = await client.send(req, stream=True)
-        
-        ctype = r.headers.get('content-type', '').lower()
-        if 'text/html' in ctype:
-            await r.aclose()
-            raise HTTPException(status_code=415, detail="Invalid Media Type: HTML Received")
-            
-        return await process_response(r, filename)
-    except HTTPException:
-        raise
+        return await process_response(r)
     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, _ = 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'
-
-    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}
-    
-    res_headers['Content-Type'] = mime_type
-    res_headers['Accept-Ranges'] = 'bytes'
+async def process_response(r):
+    excluded = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
+    response_headers = {n: v for n, v in r.headers.items() if n.lower() not in excluded}
     
-    quoted_name = urllib.parse.quote(filename)
-    res_headers['Content-Disposition'] = f'inline; filename="{quoted_name}"'
-
-    async def stream_generator():
-        try:
-            async for chunk in r.aiter_bytes(chunk_size=262144):
-                yield chunk
-        finally:
-            await r.aclose()
+    # Video Seeking အတွက် မရှိမဖြစ်လိုအပ်သော Headers များ
+    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']
 
     return StreamingResponse(
-        stream_generator(),
+        r.aiter_bytes(chunk_size=131072), # 128KB chunks for perfect buffering
         status_code=r.status_code,
-        headers=res_headers,
-        media_type=mime_type
+        headers=response_headers,
+        media_type=r.headers.get('Content-Type', 'application/octet-stream')
     )
 
 if __name__ == "__main__":