Update main.py

main.py

@@ -28,6 +28,7 @@ ua = UserAgent(fallback='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/5
 MEDIAFIRE_CACHE = {}
 CACHE_TTL = 3600 
 
+# Global Client
 client = httpx.AsyncClient(
     timeout=httpx.Timeout(60.0, read=None), 
     follow_redirects=True,
@@ -79,20 +80,17 @@ async def download_proxy(request: Request, url: str, key: str = None):
     # --- MediaFire Section ---
     if "mediafire.com" in clean_url:
         target_link = None
-        
-        # 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:
                 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("'", "")
@@ -104,13 +102,12 @@ async def download_proxy(request: Request, url: str, key: str = None):
                         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"Scraper Error: {e}")
+            except: pass
 
         if target_link:
             return await stream_file(target_link, range_header, filename)
         else:
-            raise HTTPException(status_code=404, detail="MediaFire link could not be resolved.")
+            raise HTTPException(status_code=404, detail="MediaFire direct link not found.")
 
     # --- Google Drive Section ---
     elif "drive.google.com" in clean_url:
@@ -126,37 +123,60 @@ async def download_proxy(request: Request, url: str, key: str = None):
             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
+            
+            # Google Drive API Stream
             try:
-                # 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()
+                return await stream_file(api_link, range_header, filename, extra_headers=headers)
             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)
 
     else:
         return await stream_file(clean_url, range_header, filename)
 
-async def stream_file(target_url, range_header, filename):
+async def stream_file(target_url, range_header, filename, extra_headers=None):
     headers = {'User-Agent': ua.random}
     if range_header: headers['Range'] = range_header
+    if extra_headers: headers.update(extra_headers)
     
+    # Correct Way to Stream with HTTPX
     try:
-        # Request with stream=True
-        r = await client.get(target_url, headers=headers, stream=True)
+        # We use a context manager to ensure the request is properly handled
+        async def response_streamer():
+            async with client.stream("GET", target_url, headers=headers) as r:
+                if r.status_code >= 400:
+                    yield b"Error accessing file"
+                    return
+
+                mime_type, _ = mimetypes.guess_type(filename)
+                if not mime_type or 'application' in mime_type:
+                    mime_type = 'video/mp4' if filename.lower().endswith('.mp4') else 'video/x-matroska'
+
+                # Headers for Resume/Seek
+                res_headers = {
+                    'Content-Type': mime_type,
+                    'Accept-Ranges': 'bytes',
+                    'Content-Disposition': f'inline; filename="{urllib.parse.quote(filename)}"',
+                    'Cache-Control': 'no-cache'
+                }
+                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']
+
+                # Note: StreamingResponse requires the generator to be passed, 
+                # but we need to send headers first. So we use a nested approach.
+                # Since FastAPI's StreamingResponse needs an async iterator:
+                
+        # ပြန်လည်ပြင်ဆင်��ားသော Stream Logic
+        req = client.build_request("GET", target_url, headers=headers)
+        r = await client.send(req, stream=True)
         
-        # If MediaFire returns HTML (Error), raise exception to retry or show error
-        if "text/html" in r.headers.get("Content-Type", "").lower():
+        if "text/html" in r.headers.get("Content-Type", "").lower() and r.status_code == 200:
             await r.aclose()
-            raise HTTPException(status_code=415, detail="MediaFire blocked the request. Try again.")
+            raise HTTPException(status_code=415, detail="MediaFire Blocked (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))
@@ -166,22 +186,18 @@ async def process_response(r, filename):
     if not mime_type or 'application' in mime_type:
         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'
     }
-    
-    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']
+    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=131072): # 128KB chunks
+            async for chunk in r.aiter_bytes(chunk_size=131072):
                 yield chunk
         finally:
             await r.aclose()