Update main.py

main.py

@@ -8,7 +8,7 @@ from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import StreamingResponse
 import uvicorn
 
-# Google Auth Libraries များကို Import လုပ်ခြင်း
+# Google Auth Libraries
 from google.oauth2 import service_account
 import google.auth.transport.requests
 
@@ -16,7 +16,6 @@ app = FastAPI()
 
 # --- Configurations ---
 ACCESS_KEY = os.getenv("ACCESS_KEY", "786969")
-# Service Account ရဲ့ JSON ဖိုင်ထဲက စာသားအကုန်လုံးကို Copy ကူးပြီး Secret ထဲမှာ ထည့်ရပါမည်
 SERVICE_ACCOUNT_JSON_STR = os.getenv("GOOGLE_SERVICE_ACCOUNT_JSON")
 
 @app.get("/")
@@ -24,97 +23,98 @@ def index():
     return {"message": "Proxy is Online!", "usage": "/download?url=[LINK]&key=[YOUR_KEY]"}
 
 def get_google_access_token():
-    """Service Account JSON ကို အသုံးပြု၍ Access Token ရယူခြင်း"""
     if not SERVICE_ACCOUNT_JSON_STR:
-        print("Missing GOOGLE_SERVICE_ACCOUNT_JSON credential!")
         return None
-    
     try:
-        # Secret ထဲမှ JSON string ကို Dictionary အဖြစ် ပြောင်းလဲခြင်း
         creds_dict = json.loads(SERVICE_ACCOUNT_JSON_STR)
-        
-        # Scopes သတ်မှတ်ခြင်း (Drive ကို ဖတ်ရန်)
         scopes = ['https://www.googleapis.com/auth/drive.readonly']
-        
-        # Service Account Credentials တည်ဆောက်ခြင်း
-        creds = service_account.Credentials.from_service_account_info(
-            creds_dict,
-            scopes=scopes
-        )
-        
-        # Token အသစ်တောင်းခံခြင်း
+        creds = service_account.Credentials.from_service_account_info(creds_dict, scopes=scopes)
         auth_req = google.auth.transport.requests.Request()
         creds.refresh(auth_req)
-        
         return creds.token
     except Exception as e:
         print(f"Service Account Auth Error: {e}")
         return None
 
 def get_direct_url(url):
+    """Link အမျိုးအစားအလိုက် URL၊ Token နှင့် File ID တို့ကို ပြန်ပေးရန်"""
     if "drive.google.com" in url:
         fid = re.search(r'/(?:d|file/d|open\?id=)/([a-zA-Z0-9_-]+)', url)
         if not fid:
-            return None, None
+            return None, None, None
         file_id = fid.group(1)
         
         access_token = get_google_access_token()
-        if access_token:
-            # Service Account မှ Token ရပါက API v3 Link ဖြင့် ပြန်ပေးမည်
-            return f"https://www.googleapis.com/drive/v3/files/{file_id}?alt=media", access_token
-        
-        return None, None
+        # ပထမဦးစွာ Service Account Link ကို ပြန်ပေးမည်
+        target = f"https://www.googleapis.com/drive/v3/files/{file_id}?alt=media"
+        return target, access_token, file_id
         
     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, None
-        except: 
-            return None, None
+            return btn.get('href'), None, None
+        except: return None, None, None
         
     elif "dropbox.com" in url:
-        return url.replace("?dl=0", "").split("?")[0] + "?dl=1", None
+        return url.replace("?dl=0", "").split("?")[0] + "?dl=1", None, 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", None
+        return f"https://api.onedrive.com/v1.0/shares/u!{encoded_url}/root/content", None, None
         
-    return None, None
+    return None, 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, token = get_direct_url(url)
+    target_link, token, file_id = get_direct_url(url)
     if not target_link:
-        raise HTTPException(status_code=400, detail="Invalid Link or Auth Failed")
+        raise HTTPException(status_code=400, detail="Invalid Link")
 
-    # Header ပြင်ဆင်ခြင်း
     headers = {}
     range_header = request.headers.get('range')
     if range_header:
         headers['Range'] = range_header
     
-    # Token ရှိပါက Bearer Token ထည့်ပေးခြင်း
     if token:
         headers['Authorization'] = f"Bearer {token}"
 
     session = requests.Session()
-    try:
-        r = session.get(target_link, headers=headers, stream=True, allow_redirects=True)
-        r.raise_for_status() # 403 / 404 ဖြစ်ပါက တန်းပြီး Error ဖမ်းမည်
-    except requests.exceptions.RequestException as e:
-        raise HTTPException(status_code=400, detail=f"Failed to fetch file: {str(e)}")
+    
+    # ၁။ Service Account ဖြင့် အရင်ကြိုးစားကြည့်မည်
+    r = session.get(target_link, headers=headers, stream=True, allow_redirects=True)
+    
+    # ၂။ အကယ်၍ 404 တက်ခဲ့လျှင် (Public File ဖြစ်နိုင်သဖြင့်) Fallback လုပ်မည်
+    if r.status_code == 404 and "drive.google.com" in url and file_id:
+        print(f"SA 404. Trying Public Fallback for: {file_id}")
         
-    # Sign-in page သို့မဟုတ် HTML Error Page တက်လာပါက
-    if "text/html" in r.headers.get("Content-Type", "").lower() and "google.com" in target_link:
-        raise HTTPException(status_code=401, detail="Google Authentication Failed. Token may be invalid.")
+        # Public Link သုံးမည်ဖြစ်၍ Authorization ကို ဖယ်ထုတ်ရပါမည်
+        public_headers = {k: v for k, v in headers.items() if k.lower() != 'authorization'}
+        public_url = f"https://drive.google.com/uc?export=download&id={file_id}"
+        
+        # Public Link ကို စမ်းခေါ်ခြင်း
+        r = session.get(public_url, headers=public_headers, stream=True, allow_redirects=True)
+        
+        # GDrive 100MB+ Virus Scan Confirm ကျော်ရန် Logic
+        if "text/html" in r.headers.get('Content-Type', '').lower():
+            confirm_token = None
+            for k, v in session.cookies.items():
+                if k.startswith("download_warning"):
+                    confirm_token = v
+                    break
+            
+            if confirm_token:
+                r = session.get(f"{public_url}&confirm={confirm_token}", headers=public_headers, stream=True)
+
+    # Error ဆက်ရှိနေသေးပါက
+    if r.status_code >= 400:
+        raise HTTPException(status_code=r.status_code, detail=f"Failed to fetch file: {r.status_code}")
 
-    # Header များကို ပြန်လည်ပေးပို့ရန် စစ်ထုတ်ခြင်း
+    # Header ပြန်ပို့ရန် ပြင်ဆင်ခြင်း
     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'