Update main.py

main.py

@@ -1,20 +1,16 @@
-from fastapi import FastAPI, HTTPException, Query, Request
-from fastapi.responses import StreamingResponse
-import requests
-from bs4 import BeautifulSoup
+import os
 import re
 import base64
-import os
-from fastapi import FastAPI, HTTPException, Query, Request
+import requests
+from bs4 import BeautifulSoup
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import StreamingResponse
+import uvicorn
 
 app = FastAPI()
 
-@app.get("/")
-def index():
-    return {"message": "Proxy is Online!", "usage": "/download?url=[LINK]&key=[YOUR_KEY]"}
-
 # --- Configurations ---
-ACCESS_KEY = os.getenv("ACCESS_KEY", "786969") # တကယ်လို့ Secret မရှိရင် default သုံးမယ်
+ACCESS_KEY = os.getenv("ACCESS_KEY", "786969") # Secret မရှိလျှင် default အဖြစ် "786969" ကို သုံးမည်
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")
 
 @app.get("/")
@@ -26,7 +22,7 @@ def get_direct_url(url):
         fid = re.search(r'/(?:d|file/d|open\?id=)/([a-zA-Z0-9_-]+)', url)
         if fid:
             file_id = fid.group(1)
-            # API Key ထည့်သွင်းထားပါက API ကို အသုံးပြုမည် (Virus warning ကို အလိုလိုကျော်ဖြတ်ပြီးသားဖြစ်မည်)
+            # API Key ထည့်သွင်းထားပါက API ကို အသုံးပြုမည်
             if GOOGLE_API_KEY and GOOGLE_API_KEY != "YOUR_GOOGLE_API_KEY_HERE":
                 return f"https://www.googleapis.com/drive/v3/files/{file_id}?alt=media&key={GOOGLE_API_KEY}"
             else:
@@ -40,7 +36,8 @@ def get_direct_url(url):
             soup = BeautifulSoup(r.text, 'html.parser')
             btn = soup.find('a', {'id': 'downloadButton'})
             return btn.get('href') if btn else None
-        except: return None
+        except: 
+            return None
         
     elif "dropbox.com" in url:
         return url.replace("?dl=0", "").split("?")[0] + "?dl=1"
@@ -60,7 +57,7 @@ async def download_proxy(request: Request, url: str, key: str = None):
     if not target_link:
         raise HTTPException(status_code=400, detail="Invalid Link")
 
-    # Client ဆီကလာတဲ့ Range Header ကို ရယူခြင်း (Video ကျော်ကြည့်ရန်အတွက်)
+    # Client ဆီကလာတဲ့ Range Header ကို ရယူခြင်း (Video ကျော်ကြည့်ရန်)
     range_header = request.headers.get('range')
     headers = {}
     if range_header:
@@ -74,21 +71,24 @@ async def download_proxy(request: Request, url: str, key: str = None):
     
     # ---------------------------------------------------------
     # GDrive 100MB အထက် ဖိုင်များအတွက် Virus Scan Warning ကို ကျော်ဖြတ်ခြင်း
-    # (API Key မသုံးထားသော ရိုးရိုး Link များအတွက် အထူးလိုအပ်ပါသည်)
     # ---------------------------------------------------------
     if "drive.google.com" in url and (not GOOGLE_API_KEY or GOOGLE_API_KEY == "YOUR_GOOGLE_API_KEY_HERE"):
-        if "Google Drive - Virus scan warning" in r.text or "confirm=" not in target_link:
-            confirm_token = None
-            # Cookie မှတစ်ဆင့် Virus Scan Token ကို ရှာဖွေခြင်း
-            for cookie_key, cookie_value in session.cookies.items():
-                if cookie_key.startswith("download_warning"):
-                    confirm_token = cookie_value
-                    break
-            
-            if confirm_token:
-                # Token ရပါက Link အသစ်ဖြင့် ဒေါင်းလုဒ် ပြန်စပါမည်
-                new_target_link = f"{target_link}&confirm={confirm_token}"
-                r = session.get(new_target_link, headers=headers, stream=True, allow_redirects=True)
+        # Content-Type သည် HTML ဖြစ်နေမှသာ (Virus Warning Page ဖြစ်နိုင်ခြေရှိမှသာ) r.text ကို ဖတ်မည်
+        content_type = r.headers.get('Content-Type', '').lower()
+        if 'text/html' in content_type:
+            text_content = r.text
+            if "Google Drive - Virus scan warning" in text_content or "confirm=" not in target_link:
+                confirm_token = None
+                # Cookie မှတစ်ဆင့် Virus Scan Token ကို ရှာဖွေခြင်း
+                for cookie_key, cookie_value in session.cookies.items():
+                    if cookie_key.startswith("download_warning"):
+                        confirm_token = cookie_value
+                        break
+                
+                if confirm_token:
+                    # Token ရပါက Link အသစ်ဖြင့် ဒေါင်းလုဒ် ပြန်စပါမည်
+                    new_target_link = f"{target_link}&confirm={confirm_token}"
+                    r = session.get(new_target_link, headers=headers, stream=True, allow_redirects=True)
 
     # Header များကို ပြန်လည်ပေးပို့ရန် ပြင်ဆင်ခြင်း
     excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
@@ -104,8 +104,9 @@ async def download_proxy(request: Request, url: str, key: str = None):
         response_headers['Content-Range'] = r.headers['Content-Range']
 
     def iterfile():
-        for chunk in r.iter_content(chunk_size=1024*1024):
-            if chunk: yield chunk
+        for chunk in r.iter_content(chunk_size=1024*1024): # 1MB chunks
+            if chunk: 
+                yield chunk
 
     return StreamingResponse(
         iterfile(),
@@ -115,5 +116,4 @@ async def download_proxy(request: Request, url: str, key: str = None):
     )
 
 if __name__ == "__main__":
-    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)