Update main.py

main.py

@@ -10,103 +10,126 @@ import uvicorn
 app = FastAPI()
 
 # --- Configurations ---
-ACCESS_KEY = os.getenv("ACCESS_KEY", "786969") # Secret မရှိလျှင် default အဖြစ် "786969" ကို သုံးမည်
+ACCESS_KEY = os.getenv("ACCESS_KEY", "786969")
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")
 
+# OAuth Credentials (Hugging Face Secrets ထဲမှာ ထည့်ထားပေးပါ)
+CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
+CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET")
+REFRESH_TOKEN = os.getenv("GOOGLE_REFRESH_TOKEN")
+
 @app.get("/")
 def index():
     return {"message": "Proxy is Online!", "usage": "/download?url=[LINK]&key=[YOUR_KEY]"}
 
+def get_google_access_token():
+    """Refresh Token ကိုသုံးပြီး Access Token အသစ်တောင်းယူခြင်း"""
+    if not (CLIENT_ID and CLIENT_SECRET and REFRESH_TOKEN):
+        return None
+    try:
+        url = "https://oauth2.googleapis.com/token"
+        data = {
+            "client_id": CLIENT_ID,
+            "client_secret": CLIENT_SECRET,
+            "refresh_token": REFRESH_TOKEN,
+            "grant_type": "refresh_token",
+        }
+        r = requests.post(url, data=data, timeout=10)
+        return r.json().get("access_token")
+    except:
+        return None
+
 def get_direct_url(url):
+    """Link အမျိုးအစားအလိုက် Direct Link နှင့် လိုအပ်ပါက Bearer Token ပြန်ပေးခြင်း"""
+    token = None
+    
     if "drive.google.com" in 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 ကို အသုံးပြုမည်
+            # OAuth သုံးနိုင်လျှင် (Private Files အတွက် အကောင်းဆုံး)
+            access_token = get_google_access_token()
+            if access_token:
+                return f"https://www.googleapis.com/drive/v3/files/{file_id}?alt=media", access_token
+            
+            # API Key ရှိလျှင် (Public Files အတွက်)
             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
-        
+                return f"https://www.googleapis.com/drive/v3/files/{file_id}?alt=media&key={GOOGLE_API_KEY}", None
+            
+            # API Key မရှိလျှင် ရိုးရိုး Direct Link
+            return f"https://drive.google.com/uc?export=download&id={file_id}", None
+            
     elif "mediafire.com" in url:
         try:
             r = requests.get(url, headers={'User-Agent': 'Mozilla/5.0'}, timeout=10)
             soup = BeautifulSoup(r.text, 'html.parser')
             btn = soup.find('a', {'id': 'downloadButton'})
-            return btn.get('href') if btn else None
-        except: 
-            return None
+            return btn.get('href') if btn else None, None
+        except: return None, None
         
     elif "dropbox.com" in url:
-        return url.replace("?dl=0", "").split("?")[0] + "?dl=1"
+        return url.replace("?dl=0", "").split("?")[0] + "?dl=1", None
         
     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 f"https://api.onedrive.com/v1.0/shares/u!{encoded_url}/root/content", None
         
-    return None
+    return None, None
 
 @app.get("/download")
 async def download_proxy(request: Request, url: str, key: str = None):
+    # Proxy Password စစ်ဆေးခြင်း
     if key != ACCESS_KEY:
         raise HTTPException(status_code=403, detail="Access Denied")
 
-    target_link = get_direct_url(url)
+    # Direct Link နှင့် Token ရယူခြင်း
+    target_link, token = get_direct_url(url)
     if not target_link:
         raise HTTPException(status_code=400, detail="Invalid Link")
 
-    # Client ဆီကလာတဲ့ Range Header ကို ရယူခြင်း (Video ကျော်ကြည့်ရန်)
-    range_header = request.headers.get('range')
+    # Header များ ပြင်ဆင်ခြင်း
     headers = {}
+    range_header = request.headers.get('range')
     if range_header:
         headers['Range'] = range_header
+    
+    # OAuth Token ရှိလျှင် Authorization Header ထည့်မည်
+    if token:
+        headers['Authorization'] = f"Bearer {token}"
 
-    # 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 ကို ကျော်ဖြတ်ခြင်း
+    # GDrive Virus Scan Bypass (OAuth မသုံးထားသော Public Links များအတွက်)
     # ---------------------------------------------------------
-    if "drive.google.com" in url and (not GOOGLE_API_KEY or GOOGLE_API_KEY == "YOUR_GOOGLE_API_KEY_HERE"):
-        # Content-Type သည် HTML ဖြစ်နေမှသာ (Virus Warning Page ဖြစ်နိုင်ခြေရှိမှသာ) r.text ကို ဖတ်မည်
+    if "drive.google.com" in url and not token:
         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:
+            # HTML ဖြစ်မှသာ r.text ကိုဖတ်မည် (RAM မပြည့်စေရန်)
+            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
+                for c_key, c_val in session.cookies.items():
+                    if c_key.startswith("download_warning"):
+                        confirm_token = c_val
                         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)
+                    new_link = f"{target_link}&confirm={confirm_token}"
+                    r = session.get(new_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 နှင့် Content-Range ထည့်ပေးရပါမည်
+    # Header များကို Proxy လုပ်၍ ပြန်ပို့ခြင်း
+    excluded = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
+    response_headers = {n: v for n, v in r.headers.items() if n.lower() not in excluded}
     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): # 1MB chunks
-            if chunk: 
-                yield chunk
+        for chunk in r.iter_content(chunk_size=1024*1024):
+            if chunk: yield chunk
 
     return StreamingResponse(
         iterfile(),