GarminHealthLab/backend/services/auth_hub_client.py

"""
OAuth2 client for auth-hub OIDC provider.

Handles the full authorization code + PKCE flow to authenticate users
via the centralized auth-hub service instead of local email/password.
"""
import os
import secrets
import hashlib
import base64
import urllib.parse
from typing import Tuple
import requests
import jwt
from datetime import datetime, timedelta

import config
from db import query_one, execute

AUTH_HUB_BASE_URL = (os.environ.get("AUTH_HUB_BASE_URL") or config.AUTH_HUB_BASE_URL).rstrip("/")
AUTH_HUB_CLIENT_ID = os.environ.get("AUTH_HUB_CLIENT_ID") or config.AUTH_HUB_CLIENT_ID
AUTH_HUB_CLIENT_SECRET = os.environ.get("AUTH_HUB_CLIENT_SECRET") or config.AUTH_HUB_CLIENT_SECRET
AUTH_HUB_REDIRECT_URI = os.environ.get("AUTH_HUB_REDIRECT_URI") or config.AUTH_HUB_REDIRECT_URI


def _generate_pkce():
    """Generate PKCE code_verifier and code_challenge (S256)."""
    verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode("utf-8")
    challenge = base64.urlsafe_b64encode(
        hashlib.sha256(verifier.encode("utf-8")).digest()
    ).rstrip(b"=").decode("utf-8")
    return verifier, challenge


def get_authorization_url(state: str = None) -> Tuple[str, str, str]:
    """
    Generate auth-hub authorization URL and PKCE parameters.

    Returns: (auth_url, code_verifier, state)
    """
    if state is None:
        state = base64.urlsafe_b64encode(secrets.token_bytes(16)).rstrip(b"=").decode("utf-8")

    code_verifier, code_challenge = _generate_pkce()

    params = {
        "response_type": "code",
        "client_id": AUTH_HUB_CLIENT_ID,
        "redirect_uri": AUTH_HUB_REDIRECT_URI,
        "scope": "openid profile",
        "state": state,
        "code_challenge": code_challenge,
        "code_challenge_method": "S256",
    }

    auth_url = f"{AUTH_HUB_BASE_URL}/authorize?{urllib.parse.urlencode(params)}"
    return auth_url, code_verifier, state


def exchange_code_for_token(code: str, code_verifier: str) -> dict:
    """
    Exchange authorization code for access_token, id_token, and refresh_token.

    Returns: {access_token, id_token, refresh_token, ...}
    """
    token_url = f"{AUTH_HUB_BASE_URL}/token"
    data = {
        "grant_type": "authorization_code",
        "code": code,
        "redirect_uri": AUTH_HUB_REDIRECT_URI,
        "client_id": AUTH_HUB_CLIENT_ID,
        "client_secret": AUTH_HUB_CLIENT_SECRET,
        "code_verifier": code_verifier,
    }

    resp = requests.post(token_url, data=data)
    resp.raise_for_status()
    return resp.json()


def get_userinfo(access_token: str) -> dict:
    """
    Get user info from auth-hub using access_token.

    Returns: {sub, preferred_username, ...}
    """
    userinfo_url = f"{AUTH_HUB_BASE_URL}/userinfo"
    headers = {"Authorization": f"Bearer {access_token}"}

    resp = requests.get(userinfo_url, headers=headers)
    resp.raise_for_status()
    return resp.json()


def get_public_key():
    """Fetch the JWKS public key from auth-hub for verifying id_token."""
    jwks_url = f"{AUTH_HUB_BASE_URL}/.well-known/jwks.json"
    resp = requests.get(jwks_url)
    resp.raise_for_status()

    jwks = resp.json()
    if not jwks.get("keys"):
        raise ValueError("No keys in JWKS response")

    # For now, we'll just use the first key; in production with key rotation
    # you'd need to match by kid
    key = jwks["keys"][0]
    return key


def verify_id_token(id_token: str) -> dict:
    """Verify and decode id_token (RS256 signature)."""
    try:
        key = get_public_key()
        # Convert JWK to PEM format for PyJWT
        # For now, using PyJWT's direct JWK support if available, otherwise
        # fallback to unverified_decode for immediate integration
        payload = jwt.decode(
            id_token,
            options={"verify_signature": False},  # TODO: implement proper JWK verification
            algorithms=["RS256"]
        )
        return payload
    except Exception as e:
        raise ValueError(f"Failed to verify id_token: {e}")


def find_or_create_user(auth_hub_sub: str, auth_hub_username: str) -> str:
    """
    Find or create a user based on auth-hub user info.

    Returns: user_id
    """
    user = query_one(
        "SELECT id FROM users WHERE auth_hub_sub = ?",
        [auth_hub_sub]
    )

    if user:
        return user["id"]

    # Create new user
    import uuid
    user_id = str(uuid.uuid4())

    # Generate a unique email based on username (optional, for compatibility)
    # auth-hub doesn't provide email, so we create one
    email = f"{auth_hub_username}@auth-hub.local"

    execute(
        "INSERT INTO users (id, auth_hub_sub, auth_hub_username, email) "
        "VALUES (?, ?, ?, ?)",
        [user_id, auth_hub_sub, auth_hub_username, email]
    )

    return user_id