Spaces:
Sleeping
Sleeping
| """Disposable OAuth scope validation app for RL-LLM-Wiki onboarding.""" | |
| from __future__ import annotations | |
| import asyncio | |
| import logging | |
| import os | |
| import re | |
| import secrets | |
| import time | |
| from typing import Any | |
| from urllib.parse import urlencode | |
| import httpx | |
| from fastapi import Cookie, FastAPI, HTTPException, Request | |
| from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse | |
| from huggingface_hub import ( | |
| CommitOperationAdd, | |
| HfApi, | |
| batch_bucket_files, | |
| create_bucket, | |
| delete_bucket, | |
| whoami, | |
| ) | |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") | |
| log = logging.getLogger("oauth-scope-probe") | |
| HUB = "https://huggingface.co" | |
| ORG = "rl-llm-wiki" | |
| ORG_ID = os.environ.get("RL_LLM_WIKI_ORG_ID", "6a3d324229e8338e21935fc0") | |
| PERSONAL_REPO = os.environ.get("PROBE_PERSONAL_REPO", "thomwolf/rl-oauth-scope-probe") | |
| OAUTH_CLIENT_ID = os.environ.get("OAUTH_CLIENT_ID", "") | |
| OAUTH_CLIENT_SECRET = os.environ.get("OAUTH_CLIENT_SECRET", "") | |
| SPACE_HOST = os.environ.get("SPACE_HOST", "") | |
| SCOPES = "openid profile email contribute-repos write-discussions" | |
| _states: dict[str, float] = {} | |
| _sessions: dict[str, dict[str, Any]] = {} | |
| app = FastAPI(title="RL Wiki OAuth scope probe") | |
| def _redirect_uri(request: Request) -> str: | |
| if SPACE_HOST: | |
| return f"https://{SPACE_HOST}/oauth/callback" | |
| return str(request.url_for("oauth_callback")) | |
| def _prune() -> None: | |
| cutoff = time.time() - 3600 | |
| for state, created_at in list(_states.items()): | |
| if created_at < cutoff: | |
| _states.pop(state, None) | |
| for sid, session in list(_sessions.items()): | |
| if float(session.get("created_at", 0)) < cutoff: | |
| _sessions.pop(sid, None) | |
| def index() -> str: | |
| return """<!doctype html> | |
| <html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1"> | |
| <title>RL Wiki OAuth scope probe</title><style> | |
| body{font:16px system-ui;max-width:760px;margin:48px auto;padding:0 20px;color:#171717} | |
| button,a{display:inline-block;padding:10px 15px;border-radius:7px;border:1px solid #3156a3;background:#3156a3;color:white;text-decoration:none;cursor:pointer} | |
| button:disabled{opacity:.5} pre{white-space:pre-wrap;background:#f5f6f8;padding:16px;border-radius:8px;line-height:1.5} | |
| </style></head><body><h1>OAuth scope probe</h1> | |
| <p>Scopes: <code>contribute-repos write-discussions email</code>.</p> | |
| <a href="/oauth/login" id="login">Sign in with Hugging Face</a> | |
| <button id="run" hidden>Run lifecycle probe</button><pre id="out">Checking session…</pre> | |
| <script> | |
| const out=document.getElementById('out'), login=document.getElementById('login'), run=document.getElementById('run'); | |
| async function status(){const r=await fetch('/probe/status');const d=await r.json(); | |
| if(d.authenticated){login.hidden=true;run.hidden=false;out.textContent='Signed in as '+d.username+'. Ready.'} | |
| else{out.textContent='Sign in to mint the target-scope token.'}} | |
| run.onclick=async()=>{run.disabled=true;out.textContent='Running…'; | |
| const r=await fetch('/probe/run',{method:'POST'});const d=await r.json();out.textContent=JSON.stringify(d,null,2);run.disabled=false;}; | |
| status();</script></body></html>""" | |
| def oauth_login(request: Request) -> RedirectResponse: | |
| if not (OAUTH_CLIENT_ID and OAUTH_CLIENT_SECRET): | |
| raise HTTPException(503, "OAuth is not configured") | |
| _prune() | |
| state = secrets.token_urlsafe(32) | |
| _states[state] = time.time() | |
| params = urlencode( | |
| { | |
| "response_type": "code", | |
| "client_id": OAUTH_CLIENT_ID, | |
| "redirect_uri": _redirect_uri(request), | |
| "scope": SCOPES, | |
| "state": state, | |
| "orgIds": ORG_ID, | |
| } | |
| ) | |
| return RedirectResponse(f"{HUB}/oauth/authorize?{params}") | |
| async def oauth_callback(request: Request) -> RedirectResponse: | |
| code = request.query_params.get("code") | |
| state = request.query_params.get("state") | |
| if not code or not state or _states.pop(state, None) is None: | |
| raise HTTPException(400, "Invalid OAuth state") | |
| async with httpx.AsyncClient(timeout=30, follow_redirects=True) as client: | |
| response = await client.post( | |
| f"{HUB}/oauth/token", | |
| auth=(OAUTH_CLIENT_ID, OAUTH_CLIENT_SECRET), | |
| data={ | |
| "grant_type": "authorization_code", | |
| "code": code, | |
| "redirect_uri": _redirect_uri(request), | |
| "client_id": OAUTH_CLIENT_ID, | |
| }, | |
| headers={"Accept": "application/json"}, | |
| ) | |
| response.raise_for_status() | |
| token_data = response.json() | |
| access_token = token_data.get("access_token") | |
| if not access_token: | |
| raise HTTPException(502, "Provider returned no access token") | |
| userinfo_response = await client.get( | |
| f"{HUB}/oauth/userinfo", | |
| headers={"Authorization": f"Bearer {access_token}"}, | |
| ) | |
| userinfo_response.raise_for_status() | |
| sid = secrets.token_urlsafe(32) | |
| _sessions[sid] = { | |
| "created_at": time.time(), | |
| "token_data": token_data, | |
| "userinfo": userinfo_response.json(), | |
| } | |
| redirect = RedirectResponse("/") | |
| redirect.set_cookie( | |
| "probe_sid", | |
| sid, | |
| max_age=3600, | |
| httponly=True, | |
| secure=True, | |
| samesite="none", | |
| ) | |
| return redirect | |
| def _session(probe_sid: str | None) -> dict[str, Any]: | |
| if not probe_sid or probe_sid not in _sessions: | |
| raise HTTPException(401, "Not signed in") | |
| return _sessions[probe_sid] | |
| def probe_status(probe_sid: str | None = Cookie(default=None)) -> dict[str, Any]: | |
| session = _sessions.get(probe_sid or "") | |
| if not session: | |
| return {"authenticated": False} | |
| userinfo = session["userinfo"] | |
| return { | |
| "authenticated": True, | |
| "username": userinfo.get("preferred_username") or userinfo.get("name") or "unknown", | |
| } | |
| async def _refresh(token_data: dict[str, Any]) -> tuple[str, dict[str, Any]]: | |
| refresh_token = token_data.get("refresh_token") | |
| report = { | |
| "present": bool(refresh_token), | |
| "succeeded": False, | |
| "access_token_rotated": None, | |
| "refresh_token_rotated": None, | |
| "expires_in": None, | |
| } | |
| if not refresh_token: | |
| return str(token_data["access_token"]), report | |
| async with httpx.AsyncClient(timeout=30) as client: | |
| response = await client.post( | |
| f"{HUB}/oauth/token", | |
| auth=(OAUTH_CLIENT_ID, OAUTH_CLIENT_SECRET), | |
| data={ | |
| "grant_type": "refresh_token", | |
| "refresh_token": refresh_token, | |
| "client_id": OAUTH_CLIENT_ID, | |
| }, | |
| headers={"Accept": "application/json"}, | |
| ) | |
| if not response.is_success: | |
| report["error"] = f"HTTP {response.status_code}" | |
| return str(token_data["access_token"]), report | |
| refreshed = response.json() | |
| new_access = str(refreshed.get("access_token") or token_data["access_token"]) | |
| new_refresh = refreshed.get("refresh_token") | |
| report.update( | |
| { | |
| "succeeded": True, | |
| "access_token_rotated": new_access != token_data["access_token"], | |
| "refresh_token_rotated": bool(new_refresh and new_refresh != refresh_token), | |
| "expires_in": refreshed.get("expires_in"), | |
| } | |
| ) | |
| return new_access, report | |
| def _run_lifecycle(token: str, userinfo: dict[str, Any]) -> dict[str, Any]: | |
| result: dict[str, Any] = { | |
| "whoami": False, | |
| "create_bucket": False, | |
| "write_bucket": False, | |
| "open_pr": False, | |
| "comment_pr": False, | |
| "close_pr": False, | |
| "personal_repo_write_denied": False, | |
| "cleanup_bucket": False, | |
| } | |
| suffix = secrets.token_hex(4) | |
| agent_id = f"test-oauth-{suffix}" | |
| bucket = f"{ORG}/rl-{agent_id}" | |
| api = HfApi(token=token) | |
| pr_num: int | None = None | |
| try: | |
| identity = whoami(token=token) | |
| username = str(identity.get("name") or "") | |
| if not username: | |
| raise RuntimeError("whoami returned no username") | |
| result["whoami"] = True | |
| result["username_matches_userinfo"] = username == ( | |
| userinfo.get("preferred_username") or userinfo.get("name") | |
| ) | |
| result["token_auth_type"] = (identity.get("auth") or {}).get("type") | |
| create_bucket(bucket, token=token) | |
| result["create_bucket"] = True | |
| batch_bucket_files( | |
| bucket, | |
| add=[(b"hello from the OAuth scope probe\n", "probe.txt")], | |
| token=token, | |
| ) | |
| result["write_bucket"] = True | |
| commit = api.create_commit( | |
| repo_id=f"{ORG}/knowledge-base", | |
| repo_type="dataset", | |
| operations=[ | |
| CommitOperationAdd( | |
| f"sources/probe-oauth-{suffix}.md", | |
| b"# OAuth scope validation probe\n", | |
| ) | |
| ], | |
| commit_message=f"source: OAuth scope validation {suffix}", | |
| commit_description=f"agent: {agent_id}\n\nTemporary OAuth scope validation probe.", | |
| create_pr=True, | |
| ) | |
| pr_num = commit.pr_num | |
| if pr_num is None: | |
| raise RuntimeError("create_commit returned no PR number") | |
| result["open_pr"] = True | |
| result["pr_num"] = pr_num | |
| api.comment_discussion( | |
| f"{ORG}/knowledge-base", | |
| pr_num, | |
| comment=f"/comment\n\nagent: {agent_id}\n\nOAuth permission probe.", | |
| repo_type="dataset", | |
| ) | |
| result["comment_pr"] = True | |
| api.change_discussion_status( | |
| f"{ORG}/knowledge-base", | |
| pr_num, | |
| new_status="closed", | |
| repo_type="dataset", | |
| ) | |
| result["close_pr"] = True | |
| try: | |
| api.create_commit( | |
| repo_id=PERSONAL_REPO, | |
| repo_type="model", | |
| operations=[CommitOperationAdd("oauth-should-fail.txt", b"denied\n")], | |
| commit_message="OAuth scope probe: this write must fail", | |
| create_pr=False, | |
| ) | |
| result["personal_repo_write_error"] = "UNEXPECTED SUCCESS" | |
| except Exception as error: | |
| result["personal_repo_write_denied"] = True | |
| result["personal_repo_denial_type"] = type(error).__name__ | |
| except Exception as error: | |
| result["error_step"] = next( | |
| (key for key, value in result.items() if value is False), "unknown" | |
| ) | |
| result["error_type"] = type(error).__name__ | |
| result["error"] = str(error)[:500] | |
| finally: | |
| if pr_num is not None and not result["close_pr"]: | |
| try: | |
| api.change_discussion_status( | |
| f"{ORG}/knowledge-base", | |
| pr_num, | |
| new_status="closed", | |
| repo_type="dataset", | |
| ) | |
| result["close_pr_cleanup"] = True | |
| except Exception as error: | |
| result["close_pr_cleanup_error"] = type(error).__name__ | |
| try: | |
| delete_bucket(bucket, missing_ok=True, token=token) | |
| result["cleanup_bucket"] = True | |
| except Exception as error: | |
| result["cleanup_bucket_error"] = type(error).__name__ | |
| required = ( | |
| "whoami", | |
| "create_bucket", | |
| "write_bucket", | |
| "open_pr", | |
| "comment_pr", | |
| "close_pr", | |
| "personal_repo_write_denied", | |
| "cleanup_bucket", | |
| ) | |
| result["passed"] = all(result[key] is True for key in required) | |
| return result | |
| async def run_probe(probe_sid: str | None = Cookie(default=None)) -> JSONResponse: | |
| session = _session(probe_sid) | |
| token, refresh_report = await _refresh(session["token_data"]) | |
| result = await asyncio.to_thread(_run_lifecycle, token, session["userinfo"]) | |
| userinfo = session["userinfo"] | |
| result["userinfo"] = { | |
| "sub_present": bool(userinfo.get("sub")), | |
| "email_present": bool(userinfo.get("email")), | |
| "organization_fields": sorted( | |
| (userinfo.get("organizations") or userinfo.get("orgs") or [{}])[0].keys() | |
| ) | |
| if (userinfo.get("organizations") or userinfo.get("orgs")) | |
| else [], | |
| } | |
| result["refresh"] = refresh_report | |
| return JSONResponse(result, headers={"Cache-Control": "no-store"}) | |