auth.py

"""
auth.py — Hugging Face token verification.

The frontend signs the user in with "Sign in with Hugging Face" and forwards the
resulting OAuth access token to this server on every scan. We verify that token
against Hugging Face's userinfo endpoint to establish *who* the caller is. The
stable identity is the `sub` field (an immutable user id); usernames can change.

This module is deliberately tiny and dependency-light (just `requests`).
"""

from __future__ import annotations

import requests

USERINFO_URL = "https://huggingface.co/oauth/userinfo"
WHOAMI_URL = "https://huggingface.co/api/whoami-v2"
_TIMEOUT_S = 10


def verify_hf_token(token: str | None) -> dict | None:
    """Return the user info dict for a valid token, else None.

    A valid response contains at least: sub, preferred_username, name, picture.

    Two token kinds are accepted:
      - OAuth access tokens (production frontend "Sign in with HF") -> verified
        via the OIDC /oauth/userinfo endpoint.
      - Personal access tokens (local testing / single-user mode) -> /oauth/userinfo
        rejects them, so we fall back to /api/whoami-v2 and map its response into
        the same shape (sub<-id, preferred_username<-name, name<-fullname,
        picture<-avatarUrl).

    Never raises — all failures collapse to None so callers can treat an invalid
    token identically to a missing one.
    """
    if not token:
        return None

    info = _verify_oauth_userinfo(token)
    if info is None:
        info = _verify_whoami(token)
    if not info or not info.get("sub"):
        return None
    return info


def _verify_oauth_userinfo(token: str) -> dict | None:
    """OIDC userinfo — works for OAuth access tokens."""
    try:
        resp = requests.get(
            USERINFO_URL,
            headers={"Authorization": f"Bearer {token}"},
            timeout=_TIMEOUT_S,
        )
    except requests.RequestException as exc:
        print(f"[auth] userinfo request failed: {exc}")
        return None
    if resp.status_code != 200:
        # 401 for an invalid/expired/non-OAuth token is expected and quiet.
        if resp.status_code != 401:
            print(f"[auth] userinfo returned {resp.status_code}")
        return None
    try:
        return resp.json()
    except ValueError:
        print("[auth] userinfo returned non-JSON body")
        return None


def _verify_whoami(token: str) -> dict | None:
    """Fallback for personal access tokens (and also accepts OAuth tokens).
    Maps the whoami-v2 response into the userinfo shape."""
    try:
        resp = requests.get(
            WHOAMI_URL,
            headers={"Authorization": f"Bearer {token}"},
            timeout=_TIMEOUT_S,
        )
    except requests.RequestException as exc:
        print(f"[auth] whoami request failed: {exc}")
        return None
    if resp.status_code != 200:
        if resp.status_code != 401:
            print(f"[auth] whoami returned {resp.status_code}")
        return None
    try:
        data = resp.json()
    except ValueError:
        print("[auth] whoami returned non-JSON body")
        return None
    if data.get("type") != "user" or not data.get("id"):
        return None
    return {
        "sub": data["id"],
        "preferred_username": data.get("name", ""),
        "name": data.get("fullname") or data.get("name", ""),
        "picture": data.get("avatarUrl", ""),
    }


def extract_bearer(auth_header: str | None) -> str | None:
    """Pull the token out of an 'Authorization: Bearer <token>' header.

    Provided for any future REST-style callers; the Gradio endpoints receive the
    token as a plain parameter and don't need this.
    """
    if not auth_header:
        return None
    parts = auth_header.split(None, 1)
    if len(parts) == 2 and parts[0].lower() == "bearer":
        return parts[1].strip()
    return None