Update app.py

app.py

@@ -1,16 +1,27 @@
 import os
 import sqlite3
+import base64
 from datetime import datetime, timedelta
 from typing import Optional
-from fastapi import FastAPI, Request, HTTPException, Cookie
+from fastapi import FastAPI, Request, Response, Cookie, HTTPException
 from fastapi.responses import HTMLResponse, RedirectResponse, JSONResponse
 from fastapi.templating import Jinja2Templates
+from fastapi.staticfiles import StaticFiles
 from huggingface_hub import whoami
 import secrets
+import httpx
 
 app = FastAPI()
 templates = Jinja2Templates(directory="templates")
 
+# OAuth configuration from HF Spaces environment
+OAUTH_CLIENT_ID = os.getenv("OAUTH_CLIENT_ID")
+OAUTH_CLIENT_SECRET = os.getenv("OAUTH_CLIENT_SECRET")
+OAUTH_SCOPES = os.getenv("OAUTH_SCOPES", "openid profile")
+SPACE_HOST = os.getenv("SPACE_HOST", "localhost:7860")
+OPENID_PROVIDER_URL = os.getenv("OPENID_PROVIDER_URL", "https://huggingface.co")
+
+# Database setup
 DB_PATH = "/data/usage.db"
 
 def init_db():
@@ -59,7 +70,7 @@ def record_session(username: str, is_pro: bool):
         )
         conn.commit()
     except sqlite3.IntegrityError:
-        # Already recorded today, that's fine
+        # Already recorded today
         pass
     finally:
         conn.close()
@@ -70,39 +81,74 @@ def can_start_session(username: str, is_pro: bool) -> tuple[bool, int, int]:
     limit = 15 if is_pro else 1
     return used < limit, used, limit
 
-def verify_token(token: str) -> Optional[dict]:
-    """Verify HF token and get user info"""
+async def exchange_code_for_token(code: str, redirect_uri: str) -> dict:
+    """Exchange OAuth code for access token"""
+    token_url = f"{OPENID_PROVIDER_URL}/oauth/token"
+    
+    # Prepare Basic Auth header
+    credentials = f"{OAUTH_CLIENT_ID}:{OAUTH_CLIENT_SECRET}"
+    b64_credentials = base64.b64encode(credentials.encode()).decode()
+    
+    headers = {
+        "Authorization": f"Basic {b64_credentials}",
+        "Content-Type": "application/x-www-form-urlencoded"
+    }
+    
+    data = {
+        "grant_type": "authorization_code",
+        "code": code,
+        "redirect_uri": redirect_uri,
+        "client_id": OAUTH_CLIENT_ID
+    }
+    
+    async with httpx.AsyncClient() as client:
+        response = await client.post(token_url, data=data, headers=headers)
+        response.raise_for_status()
+        return response.json()
+
+async def get_user_info(access_token: str) -> dict:
+    """Get user info from access token using whoami"""
     try:
-        user_info = whoami(token=token)
+        user_data = whoami(token=access_token)
         return {
-            "username": user_info.get("name"),
-            "is_pro": user_info.get("isPro", False),
-            "avatar": user_info.get("avatarUrl"),
-            "fullname": user_info.get("fullname", user_info.get("name"))
+            "username": user_data.get("name"),
+            "is_pro": user_data.get("isPro", False),
+            "avatar": user_data.get("avatarUrl"),
+            "fullname": user_data.get("fullname", user_data.get("name")),
+            "email": user_data.get("email")
         }
     except Exception as e:
-        print(f"Token verification failed: {e}")
-        return None
+        print(f"Failed to get user info: {e}")
+        raise HTTPException(status_code=401, detail="Failed to get user information")
 
 @app.get("/", response_class=HTMLResponse)
-async def home(request: Request, hf_token: Optional[str] = Cookie(None)):
+async def home(request: Request, access_token: Optional[str] = Cookie(None)):
     """Home page - check auth and show app or login"""
     
-    if not hf_token:
+    if not access_token:
+        # Not logged in - show login button
         return templates.TemplateResponse("index.html", {
             "request": request,
-            "authenticated": False
+            "authenticated": False,
+            "oauth_client_id": OAUTH_CLIENT_ID,
+            "redirect_uri": f"https://{SPACE_HOST}/oauth/callback",
+            "space_host": SPACE_HOST
         })
     
-    user_info = verify_token(hf_token)
-    if not user_info:
+    # Verify token and get user info
+    try:
+        user_info = await get_user_info(access_token)
+    except:
         # Invalid token, clear it
         response = templates.TemplateResponse("index.html", {
             "request": request,
             "authenticated": False,
+            "oauth_client_id": OAUTH_CLIENT_ID,
+            "redirect_uri": f"https://{SPACE_HOST}/oauth/callback",
+            "space_host": SPACE_HOST,
             "error": "Session expired. Please login again."
         })
-        response.delete_cookie("hf_token")
+        response.delete_cookie("access_token")
         return response
     
     # Check session limits
@@ -117,44 +163,48 @@ async def home(request: Request, hf_token: Optional[str] = Cookie(None)):
         "sessions_limit": limit
     })
 
-@app.post("/api/login")
-async def login(request: Request):
-    """Login with HF token"""
-    data = await request.json()
-    token = data.get("token", "").strip()
-    
-    if not token:
-        raise HTTPException(status_code=400, detail="Token is required")
+@app.get("/oauth/callback")
+async def oauth_callback(code: str, state: Optional[str] = None):
+    """Handle OAuth callback from Hugging Face"""
+    if not code:
+        raise HTTPException(status_code=400, detail="Missing authorization code")
     
-    user_info = verify_token(token)
-    if not user_info:
-        raise HTTPException(status_code=401, detail="Invalid token")
+    redirect_uri = f"https://{SPACE_HOST}/oauth/callback"
     
-    response = JSONResponse({
-        "success": True,
-        "user": user_info
-    })
-    
-    # Set secure cookie with token (expires in 30 days)
-    response.set_cookie(
-        key="hf_token",
-        value=token,
-        httponly=True,
-        secure=False,  # Set to True in production with HTTPS
-        samesite="lax",
-        max_age=30 * 24 * 60 * 60
-    )
-    
-    return response
+    try:
+        # Exchange code for token
+        token_data = await exchange_code_for_token(code, redirect_uri)
+        access_token = token_data.get("access_token")
+        
+        if not access_token:
+            raise HTTPException(status_code=400, detail="No access token received")
+        
+        # Redirect to home with token as cookie
+        response = RedirectResponse(url="/", status_code=302)
+        response.set_cookie(
+            key="access_token",
+            value=access_token,
+            httponly=True,
+            secure=True,
+            samesite="lax",
+            max_age=30 * 24 * 60 * 60  # 30 days
+        )
+        
+        return response
+        
+    except Exception as e:
+        print(f"OAuth callback error: {e}")
+        raise HTTPException(status_code=400, detail=f"Authentication failed: {str(e)}")
 
 @app.post("/api/start-session")
-async def start_session(hf_token: Optional[str] = Cookie(None)):
+async def start_session(access_token: Optional[str] = Cookie(None)):
     """Start a new generation session"""
-    if not hf_token:
+    if not access_token:
         raise HTTPException(status_code=401, detail="Not authenticated")
     
-    user_info = verify_token(hf_token)
-    if not user_info:
+    try:
+        user_info = await get_user_info(access_token)
+    except:
         raise HTTPException(status_code=401, detail="Invalid session")
     
     can_start, used, limit = can_start_session(user_info["username"], user_info["is_pro"])
@@ -175,13 +225,14 @@ async def start_session(hf_token: Optional[str] = Cookie(None)):
     }
 
 @app.get("/api/check-limits")
-async def check_limits(hf_token: Optional[str] = Cookie(None)):
+async def check_limits(access_token: Optional[str] = Cookie(None)):
     """Check current usage limits"""
-    if not hf_token:
+    if not access_token:
         raise HTTPException(status_code=401, detail="Not authenticated")
     
-    user_info = verify_token(hf_token)
-    if not user_info:
+    try:
+        user_info = await get_user_info(access_token)
+    except:
         raise HTTPException(status_code=401, detail="Invalid session")
     
     can_start, used, limit = can_start_session(user_info["username"], user_info["is_pro"])
@@ -197,10 +248,10 @@ async def check_limits(hf_token: Optional[str] = Cookie(None)):
 async def logout():
     """Logout user"""
     response = JSONResponse({"success": True})
-    response.delete_cookie("hf_token")
+    response.delete_cookie("access_token")
     return response
 
 @app.get("/health")
 async def health():
     """Health check endpoint"""
-    return {"status": "ok"}
+    return {"status": "ok", "oauth_enabled": bool(OAUTH_CLIENT_ID)}