auth.py

"""Sign-in with Hugging Face, and the signed session cookie behind it.

Enabled by ``hf_oauth: true`` in the Space's README metadata, which injects
``OAUTH_CLIENT_ID`` / ``OAUTH_CLIENT_SECRET`` / ``OPENID_PROVIDER_URL``. Only
``openid profile`` is requested — the app needs a username and nothing else. It
never asks for repo scopes, so signing in grants it no access to your account.

Sessions are a signed cookie rather than server state: a Space can restart at
any time and there is nowhere durable to keep sessions that isn't the traces
bucket.
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import logging
import os
import secrets
import time
import urllib.parse
import urllib.request

log = logging.getLogger("tracebin.auth")

SESSION_COOKIE = "tb_session"
STATE_COOKIE = "tb_oauth"
SESSION_TTL = int(os.environ.get("SESSION_TTL_SECONDS", 8 * 3600))

# Rotating this on restart just means signing in again — one redirect, usually
# invisible. Set SESSION_SECRET to keep sessions across restarts.
SESSION_SECRET = os.environ.get("SESSION_SECRET") or secrets.token_hex(32)

_provider_cache: dict[str, str] = {}


def oauth_enabled() -> bool:
    return bool(os.environ.get("OAUTH_CLIENT_ID") and os.environ.get("OAUTH_CLIENT_SECRET"))


# ---------------------------------------------------------------------------
# signed cookies
# ---------------------------------------------------------------------------


def _b64e(raw: bytes) -> str:
    return base64.urlsafe_b64encode(raw).decode().rstrip("=")


def _b64d(text: str) -> bytes:
    return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))


def sign(payload: dict) -> str:
    body = _b64e(json.dumps(payload, separators=(",", ":")).encode())
    mac = hmac.new(SESSION_SECRET.encode(), body.encode(), hashlib.sha256).digest()
    return f"{body}.{_b64e(mac)}"


def unsign(value: str | None) -> dict | None:
    """Return the payload if the signature and expiry both hold, else None."""
    if not value or "." not in value:
        return None
    body, _, mac = value.partition(".")
    expected = hmac.new(SESSION_SECRET.encode(), body.encode(), hashlib.sha256).digest()
    try:
        if not hmac.compare_digest(_b64d(mac), expected):
            return None
        payload = json.loads(_b64d(body))
    except (ValueError, TypeError):
        return None
    if not isinstance(payload, dict):
        return None
    if payload.get("exp", 0) < time.time():
        return None
    return payload


def session_username(request) -> str | None:
    """The signed-in HF username, or None."""
    payload = unsign(request.cookies.get(SESSION_COOKIE))
    if not payload:
        return None
    username = payload.get("u")
    return username if isinstance(username, str) and username else None


def make_session(username: str) -> str:
    return sign({"u": username, "exp": int(time.time()) + SESSION_TTL})


# ---------------------------------------------------------------------------
# OIDC
# ---------------------------------------------------------------------------


def _provider() -> dict[str, str]:
    """Discover the provider's endpoints once, then reuse them."""
    if _provider_cache:
        return _provider_cache
    base = os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co").rstrip("/")
    url = f"{base}/.well-known/openid-configuration"
    with urllib.request.urlopen(url, timeout=15) as response:
        config = json.loads(response.read())
    _provider_cache.update(
        {
            "authorization_endpoint": config["authorization_endpoint"],
            "token_endpoint": config["token_endpoint"],
            "userinfo_endpoint": config["userinfo_endpoint"],
        }
    )
    return _provider_cache


def authorize_url(redirect_uri: str, state: str) -> str:
    query = urllib.parse.urlencode(
        {
            "client_id": os.environ["OAUTH_CLIENT_ID"],
            "redirect_uri": redirect_uri,
            "response_type": "code",
            "scope": "openid profile",
            "state": state,
        }
    )
    return f"{_provider()['authorization_endpoint']}?{query}"


def exchange_code(code: str, redirect_uri: str) -> str:
    """Authorization code -> access token."""
    client_id = os.environ["OAUTH_CLIENT_ID"]
    client_secret = os.environ["OAUTH_CLIENT_SECRET"]
    data = urllib.parse.urlencode(
        {
            "client_id": client_id,
            "code": code,
            "grant_type": "authorization_code",
            "redirect_uri": redirect_uri,
        }
    ).encode()
    basic = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode()
    request = urllib.request.Request(
        _provider()["token_endpoint"],
        data=data,
        headers={
            "Authorization": f"Basic {basic}",
            "Content-Type": "application/x-www-form-urlencoded",
        },
    )
    with urllib.request.urlopen(request, timeout=15) as response:
        return json.loads(response.read())["access_token"]


def fetch_username(access_token: str) -> str:
    request = urllib.request.Request(
        _provider()["userinfo_endpoint"],
        headers={"Authorization": f"Bearer {access_token}"},
    )
    with urllib.request.urlopen(request, timeout=15) as response:
        info = json.loads(response.read())
    username = info.get("preferred_username") or info.get("name")
    if not username:
        raise ValueError("userinfo response carried no username")
    return username


def new_state() -> str:
    return secrets.token_urlsafe(16)


# ---------------------------------------------------------------------------
# allow lists
# ---------------------------------------------------------------------------


def normalize_users(raw) -> list[str]:
    """Accept a list or a comma/whitespace-separated string of HF usernames."""
    if raw is None:
        return []
    if isinstance(raw, str):
        parts = raw.replace(",", " ").split()
    elif isinstance(raw, (list, tuple)):
        parts = [str(p) for p in raw]
    else:
        return []
    seen, out = set(), []
    for part in parts:
        name = part.strip().lstrip("@").lower()
        # HF usernames are alphanumeric with dashes and underscores.
        if name and name not in seen and all(c.isalnum() or c in "-_." for c in name):
            seen.add(name)
            out.append(name)
    return out[:100]


def is_allowed(username: str | None, allowed: list[str] | None) -> bool:
    if not allowed:
        return True  # unrestricted: the link is the credential
    return bool(username) and username.lower() in {a.lower() for a in allowed}