Update main.py

main.py

@@ -9,7 +9,7 @@ import urllib.parse
 import mimetypes
 from bs4 import BeautifulSoup
 from fastapi import FastAPI, HTTPException, Request
-from fastapi.responses import StreamingResponse, RedirectResponse
+from fastapi.responses import StreamingResponse
 from fake_useragent import UserAgent
 import uvicorn
 
@@ -26,12 +26,12 @@ ua = UserAgent(fallback='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/5
 
 # MediaFire Cache
 MEDIAFIRE_CACHE = {}
-CACHE_TTL = 1800 
+CACHE_TTL = 3600 
 
 client = httpx.AsyncClient(
     timeout=httpx.Timeout(60.0, read=None), 
     follow_redirects=True,
-    limits=httpx.Limits(max_connections=100, max_keepalive_connections=20)
+    limits=httpx.Limits(max_connections=200, max_keepalive_connections=50)
 )
 
 @app.get("/")
@@ -63,7 +63,7 @@ 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 "video.mp4"
     return name
 
 @app.get("/download")
@@ -78,45 +78,39 @@ async def download_proxy(request: Request, url: str, key: str = None):
 
     # --- MediaFire Section ---
     if "mediafire.com" in clean_url:
-        cached_data = MEDIAFIRE_CACHE.get(clean_url)
-        
         target_link = None
-        if cached_data and (current_time - cached_data['time']) < CACHE_TTL:
-            target_link = cached_data['link']
-        else:
+        
+        # 1. Check Cache
+        cached = MEDIAFIRE_CACHE.get(clean_url)
+        if cached and (current_time - cached['time']) < CACHE_TTL:
+            target_link = cached['link']
+        
+        # 2. Fetch New Link if not cached
+        if not target_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
-                    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}"
-                        MEDIAFIRE_CACHE[clean_url] = {'link': target_link, 'time': current_time}
+                headers = {'User-Agent': ua.random}
+                async with httpx.AsyncClient(headers=headers, follow_redirects=True, timeout=30.0) as temp_client:
+                    r = await temp_client.get(clean_url)
+                    if r.status_code == 200:
+                        # Regex for direct link
+                        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("'", "")
+                        else:
+                            soup = BeautifulSoup(r.text, 'html.parser')
+                            btn = soup.find('a', {'id': 'downloadButton'}) or soup.find('a', {'class': 'input_btn_p'})
+                            if btn: target_link = btn.get('href')
+
+                        if target_link:
+                            if target_link.startswith("//"): target_link = f"https:{target_link}"
+                            MEDIAFIRE_CACHE[clean_url] = {'link': target_link, 'time': current_time}
             except Exception as e:
-                print(f"MediaFire Error: {e}")
-                raise HTTPException(status_code=500, detail=str(e))
+                print(f"Scraper Error: {e}")
 
         if target_link:
-            # FIX: Player က Video မှန်းသိအောင် Link အဆုံးမှာ /filename.mkv ကို Hint အနေနဲ့ ထည့်ပေးလိုက်ခြင်း
-            # ဒါဟာ MediaFire ဆီ Request ရောက်တဲ့အခါ query parameter အနေနဲ့ပဲ မြင်သွားမှာမို့ link မပျက်ပါဘူး
-            separator = "&" if "?" in target_link else "?"
-            hinted_link = f"{target_link}{separator}ext={filename}"
-            return RedirectResponse(url=hinted_link)
+            return await stream_file(target_link, range_header, filename)
         else:
-            raise HTTPException(status_code=404, detail="Direct link extraction failed")
+            raise HTTPException(status_code=404, detail="MediaFire link could not be resolved.")
 
     # --- Google Drive Section ---
     elif "drive.google.com" in clean_url:
@@ -133,13 +127,14 @@ async def download_proxy(request: Request, url: str, key: str = None):
             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)
+                # Use stream=True to handle large video files
+                r = await client.get(api_link, headers=headers, stream=True)
                 if r.status_code in [200, 206]:
                     return await process_response(r, filename)
                 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, filename)
 
@@ -149,34 +144,54 @@ async def download_proxy(request: Request, url: str, key: str = None):
 async def stream_file(target_url, range_header, filename):
     headers = {'User-Agent': ua.random}
     if range_header: headers['Range'] = range_header
+    
     try:
-        req = client.build_request("GET", target_url, headers=headers)
-        r = await client.send(req, stream=True)
+        # Request with stream=True
+        r = await client.get(target_url, headers=headers, stream=True)
+        
+        # If MediaFire returns HTML (Error), raise exception to retry or show error
+        if "text/html" in r.headers.get("Content-Type", "").lower():
+            await r.aclose()
+            raise HTTPException(status_code=415, detail="MediaFire blocked the request. Try again.")
+            
         return await process_response(r, filename)
+    except HTTPException:
+        raise
     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:
-        mime_type = 'video/x-matroska' if filename.endswith('.mkv') else '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'
+        mime_type = 'video/mp4' if filename.lower().endswith('.mp4') else 'video/x-matroska'
+
+    # Essential Headers for Resume and Seeking
+    res_headers = {
+        'Content-Type': mime_type,
+        'Accept-Ranges': 'bytes',
+        'Content-Disposition': f'inline; filename="{urllib.parse.quote(filename)}"',
+        'Cache-Control': 'no-cache'
+    }
     
-    quoted_name = urllib.parse.quote(filename)
-    res_headers['Content-Disposition'] = f'inline; filename="{quoted_name}"'
+    if 'content-length' in r.headers:
+        res_headers['Content-Length'] = r.headers['content-length']
+    if 'content-range' in r.headers:
+        res_headers['Content-Range'] = r.headers['content-range']
 
     async def stream_generator():
         try:
-            async for chunk in r.aiter_bytes(chunk_size=262144):
+            async for chunk in r.aiter_bytes(chunk_size=131072): # 128KB chunks
                 yield chunk
         finally:
             await r.aclose()
 
-    return StreamingResponse(stream_generator(), status_code=r.status_code, headers=res_headers, media_type=mime_type)
+    return StreamingResponse(
+        stream_generator(),
+        status_code=r.status_code,
+        headers=res_headers,
+        media_type=mime_type
+    )
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)