nas-tools/dashboard/backend/routers/passkey.py

"""WebAuthn Passkey authentication endpoints."""

import json
import logging
import time
import uuid
from datetime import timedelta
from typing import Any

from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from slowapi import Limiter
from slowapi.util import get_remote_address
from webauthn import (
    generate_authentication_options,
    generate_registration_options,
    verify_authentication_response,
    verify_registration_response,
)
from webauthn.helpers import base64url_to_bytes, bytes_to_base64url, options_to_json
from webauthn.helpers.structs import PublicKeyCredentialDescriptor, UserVerificationRequirement

import auth_service as auth
import config

router = APIRouter()
limiter = Limiter(key_func=get_remote_address)
logger = logging.getLogger(__name__)

# In-memory challenge store with TTL: {challenge_id: (challenge_bytes, timestamp)}
_challenges: dict[str, tuple[bytes, float]] = {}


def _store_challenge(challenge: bytes) -> str:
    """Store a WebAuthn challenge and return a session-bound challenge_id."""
    challenge_id = uuid.uuid4().hex
    _challenges[challenge_id] = (challenge, time.time())
    # Clean up expired challenges (>5 minutes old)
    now = time.time()
    expired = [k for k, v in _challenges.items() if now - v[1] > 300]
    for k in expired:
        _challenges.pop(k, None)
    return challenge_id


def _get_challenge(challenge_id: str | None, client_data_b64: str) -> bytes:
    """Look up challenge by session-bound challenge_id and validate clientDataJSON."""
    if not challenge_id:
        raise HTTPException(status_code=400, detail="Missing challenge_id")
    if not client_data_b64:
        raise HTTPException(status_code=400, detail="Missing clientDataJSON")
    try:
        data = json.loads(base64url_to_bytes(client_data_b64).decode("utf-8"))
        chal_bytes = base64url_to_bytes(data.get("challenge", ""))
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid clientDataJSON")

    entry = _challenges.pop(challenge_id, None)
    if not entry or time.time() - entry[1] > 300:
        raise HTTPException(status_code=400, detail="Challenge expired or invalid")
    stored_challenge = entry[0]
    if stored_challenge != chal_bytes:
        raise HTTPException(status_code=400, detail="Challenge mismatch")
    return chal_bytes


# Pydantic models for passkey request validation
class PasskeyVerifyRequest(BaseModel):
    id: str = ""
    rawId: str = ""
    response: dict[str, Any] = {}
    type: str = "public-key"
    challenge_id: str | None = None
    name: str = ""
    # Allow extra fields for authenticator-specific extensions
    model_config = {"extra": "allow"}


class PasskeyDeleteRequest(BaseModel):
    id: str


@router.post("/passkey/register/options")
@limiter.limit("5/minute")
async def passkey_register_options(request: Request, current_user=Depends(auth.get_current_user)):
    """Generate WebAuthn registration options for adding a new passkey."""
    existing = auth.load_passkey_credentials()
    exclude = [PublicKeyCredentialDescriptor(id=base64url_to_bytes(c["credential_id"])) for c in existing]
    options = generate_registration_options(
        rp_id=config.WEBAUTHN_RP_ID,
        rp_name="NAS Dashboard",
        user_id=current_user.username.encode(),
        user_name=current_user.username,
        user_display_name=current_user.username,
        exclude_credentials=exclude,
    )
    chal_id = _store_challenge(options.challenge)
    result = json.loads(options_to_json(options))
    result["challenge_id"] = chal_id
    return JSONResponse(content=result)


@router.post("/passkey/register/verify")
async def passkey_register_verify(body: PasskeyVerifyRequest, current_user=Depends(auth.get_current_user)):
    """Verify and save a new passkey registration."""
    challenge = _get_challenge(body.challenge_id, body.response.get("clientDataJSON", ""))
    try:
        verification = verify_registration_response(
            credential=body.model_dump(),
            expected_challenge=challenge,
            expected_rp_id=config.WEBAUTHN_RP_ID,
            expected_origin=config.WEBAUTHN_ORIGINS,
        )
    except Exception:
        logger.exception("Passkey registration verification failed")
        raise HTTPException(status_code=400, detail="Invalid passkey registration")

    auth.save_passkey_credential(
        {
            "credential_id": bytes_to_base64url(verification.credential_id),
            "public_key": bytes_to_base64url(verification.credential_public_key),
            "sign_count": verification.sign_count,
            "name": body.name or "Passkey",
            "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "username": current_user.username,
            "role": current_user.role,
        }
    )
    return {"ok": True}


@router.post("/passkey/login/options")
@limiter.limit("10/minute")
async def passkey_login_options(request: Request):
    """Generate WebAuthn authentication options for passkey login."""
    creds = auth.load_passkey_credentials()
    if not creds:
        raise HTTPException(status_code=404, detail="No passkeys registered")

    allow = [PublicKeyCredentialDescriptor(id=base64url_to_bytes(c["credential_id"])) for c in creds]
    options = generate_authentication_options(
        rp_id=config.WEBAUTHN_RP_ID,
        allow_credentials=allow,
        user_verification=UserVerificationRequirement.PREFERRED,
    )
    chal_id = _store_challenge(options.challenge)
    result = json.loads(options_to_json(options))
    result["challenge_id"] = chal_id
    return JSONResponse(content=result)


@router.post("/passkey/login/verify")
@limiter.limit("10/minute")
async def passkey_login_verify(body: PasskeyVerifyRequest, request: Request, response: Response):
    """Verify passkey authentication and issue tokens."""
    challenge = _get_challenge(body.challenge_id, body.response.get("clientDataJSON", ""))
    creds = auth.load_passkey_credentials()
    cred_id_b64 = body.id
    matched = next((c for c in creds if c["credential_id"] == cred_id_b64), None)
    if not matched:
        raise HTTPException(status_code=400, detail="Unknown credential")

    try:
        verification = verify_authentication_response(
            credential=body.model_dump(),
            expected_challenge=challenge,
            expected_rp_id=config.WEBAUTHN_RP_ID,
            expected_origin=config.WEBAUTHN_ORIGINS,
            credential_public_key=base64url_to_bytes(matched["public_key"]),
            credential_current_sign_count=matched["sign_count"],
        )
    except Exception:
        logger.exception("Passkey authentication verification failed")
        raise HTTPException(status_code=400, detail="Invalid passkey authentication")

    # Update sign count
    matched["sign_count"] = verification.new_sign_count
    data = auth._load_auth_data()
    data["passkey_credentials"] = creds
    auth._save_auth_data(data)

    # Use stored username and role from passkey credential
    username = matched.get("username", config.ADMIN_USER)
    role = matched.get("role", "admin")

    access_token = auth.create_access_token(
        data={"sub": username, "role": role},
        expires_delta=timedelta(minutes=config.ACCESS_TOKEN_EXPIRE_MINUTES),
    )
    refresh_token = auth.create_refresh_token(data={"sub": username, "role": role})
    auth.set_auth_cookies(response, access_token, refresh_token)
    return {
        "access_token": access_token,
        "refresh_token": refresh_token,
        "token_type": "bearer",
        "user": username,
        "role": role,
    }


@router.get("/passkey/list")
async def passkey_list(current_user=Depends(auth.get_current_user)):
    """List all registered passkeys for the current user."""
    creds = auth.load_passkey_credentials()
    return {
        "passkeys": [
            {"id": c["credential_id"], "name": c.get("name", "Passkey"), "created_at": c.get("created_at", "")}
            for c in creds
        ]
    }


@router.post("/passkey/delete")
async def passkey_delete(body: PasskeyDeleteRequest, current_user=Depends(auth.get_current_user)):
    """Delete a specific passkey by credential ID."""
    cred_id = body.id
    data = auth._load_auth_data()
    creds = data.get("passkey_credentials", [])
    data["passkey_credentials"] = [c for c in creds if c["credential_id"] != cred_id]
    auth._save_auth_data(data)
    return {"ok": True}


@router.post("/passkey/clear")
async def passkey_clear(current_user=Depends(auth.get_current_user)):
    """Clear all passkeys for the current user."""
    auth.clear_passkey_credentials()
    return {"ok": True}