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
+from fastapi.responses import StreamingResponse, RedirectResponse
 from fake_useragent import UserAgent
 import uvicorn
 
@@ -24,9 +24,9 @@ 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
 MEDIAFIRE_CACHE = {}
-CACHE_TTL = 1800 # Cache သက်တမ်း စက္ကန့် 1800 (မိနစ် 30)
+CACHE_TTL = 1800 
 
 client = httpx.AsyncClient(
     timeout=httpx.Timeout(60.0, read=None), 
@@ -74,39 +74,25 @@ async def download_proxy(request: Request, url: str, key: str = None):
     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 ---
+    # --- MediaFire Section (Modified to use Redirect) ---
     if "mediafire.com" in clean_url:
-        
-        # 1. Cache ထဲမှာရှိမရှိ အရင်စစ်ဆေးခြင်း (Player များ Seek လုပ်ရာတွင် လွယ်ကူစေရန်)
         cached_data = MEDIAFIRE_CACHE.get(clean_url)
+        
+        # Cache ရှိရင် တိုက်ရိုက် Redirect လုပ်မည်
         if cached_data and (current_time - cached_data['time']) < CACHE_TTL:
-            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 ဖြစ်သွားပါက အသစ်ပြန်ရှာခြင်း
+            return RedirectResponse(url=cached_data['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 page access failed")
+                    raise HTTPException(status_code=page_res.status_code, detail="MediaFire 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)
                 if match:
@@ -115,30 +101,25 @@ async def download_proxy(request: Request, url: str, key: str = None):
                 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 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}"
                     
-                    # နောက်ထပ် Range Request များအတွက် Memory ထဲတွင် မှတ်သားထားခြင်း
-                    MEDIAFIRE_CACHE[clean_url] = {
-                        'link': target_link,
-                        'ua': req_ua,
-                        'cookies': dict(cookies),
-                        'time': current_time
-                    }
+                    # သိမ်းဆည်းခြင်း
+                    MEDIAFIRE_CACHE[clean_url] = {'link': target_link, 'time': current_time}
                     
-                    return await stream_file(target_link, range_header, filename, referer=clean_url, req_ua=req_ua, cookies=cookies)
+                    # Proxy အစား RedirectResponse ကိုသုံးပြီး Player ဆီ လွှဲပေးလိုက်ခြင်း
+                    return RedirectResponse(url=target_link)
                 else:
-                    raise HTTPException(status_code=404, detail="Could not extract direct download link")
+                    raise HTTPException(status_code=404, detail="Direct link extraction failed")
 
         except Exception as e:
-            print(f"MediaFire Logic Error: {str(e)}")
+            print(f"MediaFire Error: {e}")
             raise HTTPException(status_code=500, detail=str(e))
 
-    # --- Google Drive Section ---
+    # --- Google Drive Section (Stay as Proxy) ---
     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")
@@ -161,42 +142,28 @@ async def download_proxy(request: Request, url: str, key: str = None):
             except: continue
 
         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)
+        return await stream_file(public_url, range_header, filename)
 
     else:
-        return await stream_file(clean_url, range_header, filename, req_ua=req_ua)
+        return await stream_file(clean_url, range_header, filename)
 
-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}
+async def stream_file(target_url, range_header, filename):
+    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
     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'
+        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'
     
@@ -210,12 +177,7 @@ async def process_response(r, filename):
         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)