Update main.py

main.py

@@ -5,6 +5,7 @@ import base64
 import httpx
 import random
 import asyncio
+import urllib.parse
 from bs4 import BeautifulSoup
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import StreamingResponse
@@ -22,16 +23,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')
 
-# Global Async Client: Video ကျော်ကြည့်တဲ့အခါ Connection အသစ်တွေ မြန်မြန်ဆောက်နိုင်ဖို့
+# Global Async Client
 client = httpx.AsyncClient(
-    timeout=httpx.Timeout(20.0, read=None), 
+    timeout=httpx.Timeout(30.0, read=None), # Timeout ကို 30 စက္ကန့်အထိ တိုးပေးထားပါတယ်
     follow_redirects=True,
     limits=httpx.Limits(max_connections=100, max_keepalive_connections=20)
 )
 
 @app.get("/")
 async def index():
-    return {"Online"}
+    return {"status": "Online"}
 
 def get_all_service_accounts():
     if not SERVICE_ACCOUNT_JSON_STR: return []
@@ -44,7 +45,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 ဖြစ်လို့)
+        # Token refresh ကို thread ထဲမှာ run ပါတယ်
         loop = asyncio.get_event_loop()
         auth_req = google.auth.transport.requests.Request()
         await loop.run_in_executor(None, creds.refresh, auth_req)
@@ -60,6 +61,9 @@ 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)
+
     range_header = request.headers.get('range')
 
     # --- MediaFire & Others ---
@@ -71,8 +75,9 @@ async def download_proxy(request: Request, url: str, key: str = None):
             try:
                 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)
+                
+                # 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:
@@ -111,7 +116,6 @@ async def download_proxy(request: Request, url: str, key: str = None):
         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]:
@@ -132,10 +136,11 @@ async def stream_file(target_url, range_header):
         r = await client.send(req, stream=True)
         return await process_response(r)
     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):
-    # မူလလာတဲ့ content-disposition ကို ဖယ်ထုတ်ဖို့ excluded ထဲ ထည့်လိုက်ပါမယ်
+    # မူလလာတဲ့ content-disposition ကို ဖယ်ထုတ်ဖို့ excluded ထဲ ထည့်ပါမယ်
     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}
     
@@ -146,15 +151,22 @@ async def process_response(r):
     if 'Content-Range' in r.headers:
         response_headers['Content-Range'] = r.headers['Content-Range']
 
-    # --- Direct Download ဖြစ်စေရန် ပြင်ဆင်ချက် ---
-    # content-disposition ကို attachment လို့ သတ်မှတ်ပေးလိုက်ရင် browser က တန်း download လုပ်ပါလိမ့်မယ်
+    # Direct Download ဖြစ်စေရန်
     response_headers['Content-Disposition'] = 'attachment'
 
+    # Connection များကို Memory Leak မဖြစ်စေဘဲ သေချာပိတ်ပေးမည့် Generator
+    async def stream_generator():
+        try:
+            async for chunk in r.aiter_bytes(chunk_size=131072): # 128KB chunks
+                yield chunk
+        finally:
+            await r.aclose()
+
     return StreamingResponse(
-        r.aiter_bytes(chunk_size=131072), # 128KB chunks for perfect buffering
+        stream_generator(),
         status_code=r.status_code,
         headers=response_headers,
-        media_type='application/octet-stream' # stream အနေနဲ့ တန်းဆွဲဖို့ application/octet-stream ပြောင်းလိုက်ပါတယ်
+        media_type='application/octet-stream'
     )
 
 if __name__ == "__main__":