proxy.py

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, parse_qs

app = FastAPI()
LT_INTERNAL = "http://127.0.0.1:5000"

CLIENT_ID     = os.environ.get("OAUTH_CLIENT_ID")
CLIENT_SECRET = os.environ.get("OAUTH_CLIENT_SECRET")
SPACE_HOST    = os.environ.get("SPACE_HOST", "localhost")
REDIRECT_URI  = f"https://{SPACE_HOST}/auth/callback"

sessions: dict[str, dict] = {}
SECURE_COOKIE = not SPACE_HOST.startswith("localhost")

# ── 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 = (
        f"https://huggingface.co/oauth/authorize"
        f"?client_id={CLIENT_ID}"
        f"&redirect_uri={quote(REDIRECT_URI)}"
        f"&response_type=code"
        f"&scope={scope}"
        f"&state={state}"
    )
    resp = RedirectResponse(url)
    resp.set_cookie("oauth_state", state, max_age=600, httponly=True, samesite="lax", secure=SECURE_COOKIE)
    return resp

# ── OAuth: callback ───────────────────────────────────────
@app.get("/auth/callback")
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(
            "https://huggingface.co/oauth/token",
            data={
                "grant_type": "authorization_code",
                "code": code,
                "redirect_uri": REDIRECT_URI,
            },
            headers={"Authorization": f"Basic {basic}"},
            timeout=30.0,
        )

    if r.status_code != 200:
        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.delete_cookie("oauth_state")
    resp.set_cookie("lt_session", session_id, httponly=True, samesite="lax", secure=SECURE_COOKIE)
    return resp

# ── 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 ""

# ── 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}"

        headers = {
            k: v for k, v in request.headers.items()
            if k.lower() not in ("host", "content-length")
        }

        try:
            lt = await client.request(
                method=request.method,
                url=url,
                headers=headers,
                content=body,
                timeout=60.0,
                follow_redirects=True,
            )
        except Exception as e:
            raise HTTPException(status_code=502, detail=f"LibreTranslate unreachable: {e}")

    excluded = {"content-encoding", "transfer-encoding", "content-length"}
    return Response(
        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":
        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."
            )

        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)