Update main.py

main.py

@@ -4,9 +4,18 @@ import requests
 from bs4 import BeautifulSoup
 import re
 import base64
+import os
+from fastapi import FastAPI, HTTPException, Query, Request
 
 app = FastAPI()
-ACCESS_KEY = "969786"
+
+@app.get("/")
+def index():
+    return {"message": "Proxy is Online!", "usage": "/download?url=[LINK]&key=[YOUR_KEY]"}
+
+# --- Configurations ---
+ACCESS_KEY = os.getenv("ACCESS_KEY", "969786") # တကယ်လို့ Secret မရှိရင် default သုံးမယ်
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")
 
 @app.get("/")
 def index():
@@ -15,7 +24,16 @@ def index():
 def get_direct_url(url):
     if "drive.google.com" in url:
         fid = re.search(r'/(?:d|file/d|open\?id=)/([a-zA-Z0-9_-]+)', url)
-        return f"https://drive.google.com/uc?export=download&id={fid.group(1)}" if fid else None
+        if fid:
+            file_id = fid.group(1)
+            # API Key ထည့်သွင်းထားပါက API ကို အသုံးပြုမည် (Virus warning ကို အလိုလိုကျော်ဖြတ်ပြီးသားဖြစ်မည်)
+            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:
+                # API Key မရှိပါက ရိုးရိုး Direct Link ကို အသုံးပြုမည်
+                return f"https://drive.google.com/uc?export=download&id={file_id}"
+        return None
+        
     elif "mediafire.com" in url:
         try:
             r = requests.get(url, headers={'User-Agent': 'Mozilla/5.0'}, timeout=10)
@@ -23,11 +41,14 @@ def get_direct_url(url):
             btn = soup.find('a', {'id': 'downloadButton'})
             return btn.get('href') if btn else None
         except: return None
+        
     elif "dropbox.com" in url:
         return url.replace("?dl=0", "").split("?")[0] + "?dl=1"
+        
     elif "1drv.ms" in url or "onedrive.live.com" in url:
         encoded_url = base64.b64encode(bytes(url, 'utf-8')).decode('utf-8').replace('=', '').replace('/', '_').replace('+', '-')
         return f"https://api.onedrive.com/v1.0/shares/u!{encoded_url}/root/content"
+        
     return None
 
 @app.get("/download")
@@ -45,17 +66,42 @@ async def download_proxy(request: Request, url: str, key: str = None):
     if range_header:
         headers['Range'] = range_header
 
-    # Target Server (MediaFire/GDrive) ဆီကို Range Header နဲ့ လှမ်းတောင်းခြင်း
-    r = requests.get(target_link, headers=headers, stream=True, allow_redirects=True)
+    # Session အသုံးပြုခြင်းဖြင့် GDrive ၏ Cookie များကို သိမ်းဆည်းနိုင်သည်
+    session = requests.Session()
+
+    # Target Server ဆီကို Range Header နဲ့ လှမ်းတောင်းခြင်း
+    r = session.get(target_link, headers=headers, stream=True, allow_redirects=True)
     
+    # ---------------------------------------------------------
+    # 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)
+
     # Header များကို ပြန်လည်ပေးပို့ရန် ပြင်ဆင်ခြင်း
     excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
     response_headers = {name: value for name, value in r.headers.items() if name.lower() not in excluded_headers}
     
-    # Video ကျော်ကြည့်နိုင်ရန် 'Accept-Ranges': 'bytes' ထည့်ပေးရပါမည်
+    # Video ကျော်ကြည့်နိုင်ရန် Accept-Ranges နှင့် Content-Range ထည့်ပေးရပါမည်
     response_headers['Accept-Ranges'] = 'bytes'
+    
     if 'Content-Length' in r.headers:
         response_headers['Content-Length'] = r.headers['Content-Length']
+        
+    if 'Content-Range' in r.headers:
+        response_headers['Content-Range'] = r.headers['Content-Range']
 
     def iterfile():
         for chunk in r.iter_content(chunk_size=1024*1024):