Update proxy.py

proxy.py

@@ -2,9 +2,10 @@ import os
 import base64
 import secrets
 import httpx
+import json
 from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import RedirectResponse, Response
-from urllib.parse import quote
+from urllib.parse import quote, parse_qs
 
 app = FastAPI()
 LT_INTERNAL = "http://127.0.0.1:5000"
@@ -14,12 +15,13 @@ CLIENT_SECRET = os.environ.get("OAUTH_CLIENT_SECRET")
 SPACE_HOST    = os.environ.get("SPACE_HOST", "localhost")
 REDIRECT_URI  = f"https://{SPACE_HOST}/auth/callback"
 
-# In-memory session store (fine for a single-container Space)
 sessions: dict[str, dict] = {}
 
-# ── OAuth login ─────────────────────────────────────────────
+# ── OAuth: login ──────────────────────────────────────────
 @app.get("/auth/login")
 async def auth_login():
+    if not CLIENT_ID:
+        raise HTTPException(status_code=500, detail="OAuth not configured")
     state = secrets.token_urlsafe(32)
     scope = quote("openid profile")
     url = (
@@ -31,13 +33,16 @@ async def auth_login():
         f"&state={state}"
     )
     resp = RedirectResponse(url)
-    resp.set_cookie("oauth_state", state, max_age=600, httponly=True, samesite="lax")
+    resp.set_cookie("oauth_state", state, max_age=600, httponly=True, samesite="lax", secure=True)
     return resp
 
-# ── OAuth callback ──────────────────────────────────────────
+# ── OAuth: callback ───────────────────────────────────────
 @app.get("/auth/callback")
-async def auth_callback(code: str, state: str):
-    # NOTE: In production, verify that `state` matches the cookie.
+async def auth_callback(code: str, state: str, request: Request):
+    cookie_state = request.cookies.get("oauth_state")
+    if not cookie_state or cookie_state != state:
+        raise HTTPException(status_code=400, detail="Invalid state parameter")
+
     basic = base64.b64encode(f"{CLIENT_ID}:{CLIENT_SECRET}".encode()).decode()
     async with httpx.AsyncClient() as client:
         r = await client.post(
@@ -50,44 +55,58 @@ async def auth_callback(code: str, state: str):
             headers={"Authorization": f"Basic {basic}"},
             timeout=30.0,
         )
+
     if r.status_code != 200:
-        raise HTTPException(status_code=400, detail="HuggingFace OAuth failed")
+        raise HTTPException(status_code=400, detail=f"HF OAuth failed: {r.text}")
 
     token_data = r.json()
     session_id = secrets.token_urlsafe(32)
     sessions[session_id] = token_data
 
     resp = RedirectResponse("/")
-    resp.set_cookie("lt_session", session_id, httponly=True, samesite="lax")
+    resp.delete_cookie("oauth_state")
+    resp.set_cookie("lt_session", session_id, httponly=True, samesite="lax", secure=True)
     return resp
 
-# ── Proxy everything ────────────────────────────────────────
-@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD", "PATCH"])
-async def proxy(request: Request, path: str):
-    # Guard /suggest behind HF OAuth
-    if path == "suggest" and request.method == "POST":
-        session = request.cookies.get("lt_session")
-        if not session or session not in sessions:
-            raise HTTPException(
-                status_code=401,
-                detail="You must be logged in with a HuggingFace account to submit suggestions. "
-                       "Visit /auth/login first."
-            )
+# ── OAuth: logout ─────────────────────────────────────────
+@app.get("/auth/logout")
+async def auth_logout():
+    resp = RedirectResponse("/")
+    resp.delete_cookie("lt_session")
+    return resp
+
+# ── Helper: extract target language from request body ─────
+def get_target_lang(body: bytes, content_type: str) -> str:
+    if not body:
+        return ""
+    if "application/json" in content_type:
+        try:
+            data = json.loads(body)
+            return data.get("target", "") or data.get("t", "")
+        except Exception:
+            return ""
+    elif "application/x-www-form-urlencoded" in content_type:
+        try:
+            form = parse_qs(body.decode("utf-8"))
+            return form.get("target", [""])[0] or form.get("t", [""])[0]
+        except Exception:
+            return ""
+    return ""
 
-    # Forward to LibreTranslate
+# ── Proxy everything to LibreTranslate ────────────────────
+async def forward(request: Request, path: str, body: bytes):
     async with httpx.AsyncClient() as client:
         url = f"{LT_INTERNAL}/{path}"
         if request.query_params:
             url = f"{url}?{request.query_params}"
 
-        body = await request.body()
         headers = {
             k: v for k, v in request.headers.items()
             if k.lower() not in ("host", "content-length")
         }
 
         try:
-            lt_resp = await client.request(
+            lt = await client.request(
                 method=request.method,
                 url=url,
                 headers=headers,
@@ -98,10 +117,35 @@ async def proxy(request: Request, path: str):
         except Exception as e:
             raise HTTPException(status_code=502, detail=f"LibreTranslate unreachable: {e}")
 
-    # Strip hop-by-hop headers
     excluded = {"content-encoding", "transfer-encoding", "content-length"}
     return Response(
-        content=lt_resp.content,
-        status_code=lt_resp.status_code,
-        headers={k: v for k, v in lt_resp.headers.items() if k.lower() not in excluded},
-    )
+        content=lt.content,
+        status_code=lt.status_code,
+        headers={k: v for k, v in lt.headers.items() if k.lower() not in excluded},
+    )
+
+@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD", "PATCH"])
+async def proxy(request: Request, path: str):
+    body = await request.body()
+    content_type = request.headers.get("content-type", "")
+
+    # ── /suggest gate: auth + kabyle-only ──────────────────
+    if path == "suggest" and request.method == "POST":
+        # 1. Must be logged in via HF
+        session = request.cookies.get("lt_session")
+        if not session or session not in sessions:
+            raise HTTPException(
+                status_code=401,
+                detail="You must be logged in with a HuggingFace account to submit suggestions. "
+                       "Visit /auth/login first."
+            )
+
+        # 2. Must be for Kabyle
+        target = get_target_lang(body, content_type)
+        if target and target != "kab":
+            raise HTTPException(
+                status_code=403,
+                detail="Suggestions are only accepted for the Kabyle (kab) language."
+            )
+
+    return await forward(request, path, body)