munch-ease-backend/api/auth_v2.py

from __future__ import annotations

import base64
import hashlib
import hmac
import os
from typing import Optional

import aiosqlite
from fastapi import APIRouter, Depends, Request
from fastapi.responses import JSONResponse

from api.deps import error_response, get_db
from common import ApiModel
from security import JwtConfig, create_jwt
from settings import settings
from users import repository as users_db

# Prefer Argon2 for new passwords; keep PBKDF2 verify for backward compatibility
try:
    from argon2 import PasswordHasher

    _ph: PasswordHasher | None = PasswordHasher()
except Exception:  # pragma: no cover - optional dependency in some environments
    _ph = None
from users.models import User

router = APIRouter(prefix="/auth", tags=["auth-v2"])


class RegisterBody(ApiModel):
    email: str
    password: str
    display_name: str


class LoginBody(ApiModel):
    email: str
    password: str


class TokenResponse(ApiModel):
    access_token: str
    token_type: str = "bearer"
    user: User


PBKDF2_ALG = "pbkdf2_sha256"
PBKDF2_ITER = 390000  # kept for verifying older hashes
SALT_BYTES = 16


def _hash_pw(pw: str) -> str:
    """Hash a password.

    Uses Argon2 when available; falls back to PBKDF2 for environments without argon2-cffi.
    """
    if _ph is not None:
        return _ph.hash(pw)
    # Fallback
    salt = os.urandom(SALT_BYTES)
    dk = hashlib.pbkdf2_hmac("sha256", pw.encode("utf-8"), salt, PBKDF2_ITER)
    return f"{PBKDF2_ALG}${PBKDF2_ITER}${base64.b64encode(salt).decode()}${base64.b64encode(dk).decode()}"


def _verify_pw(pw: str, stored: str) -> bool:
    """Verify password against either Argon2 or PBKDF2 stored hashes."""
    # Try Argon2 first
    if _ph is not None and stored.startswith("$argon2"):
        try:
            return _ph.verify(stored, pw)
        except Exception:
            return False
    # PBKDF2 fallback
    try:
        alg, iter_s, salt_b64, hash_b64 = stored.split("$", 3)
        if alg != PBKDF2_ALG:
            return False
        iters = int(iter_s)
        salt = base64.b64decode(salt_b64)
        expected = base64.b64decode(hash_b64)
        dk = hashlib.pbkdf2_hmac("sha256", pw.encode("utf-8"), salt, iters)
        return hmac.compare_digest(dk, expected)
    except Exception:
        return False


def _jwt_config() -> JwtConfig:
    # Secrets can be provided base64-encoded via env; fallback to deterministic dev defaults (NOT for prod)
    if settings.access_secret_b64:
        access = base64.b64decode(settings.access_secret_b64)
    else:
        access = b"dev-access-secret-change-me-32bytes!!"[:32]
    if settings.refresh_secret_b64:
        refresh = base64.b64decode(settings.refresh_secret_b64)
    else:
        refresh = b"dev-refresh-secret-change-me-32bytes!!"[:32]
    return JwtConfig(
        issuer=settings.jwt_issuer,
        audience=settings.jwt_audience,
        access_secret=access,
        refresh_secret=refresh,
        access_ttl_seconds=settings.access_ttl_seconds,
        refresh_ttl_seconds=settings.refresh_ttl_seconds,
    )


def _token_pair_for_user(user: User) -> tuple[str, str]:
    cfg = _jwt_config()
    sub = str(user.id)
    access = create_jwt(cfg, sub, kind="access", extra={"user_id": user.id, "email": user.email})
    refresh = create_jwt(cfg, sub, kind="refresh")
    return access, refresh


@router.post("/register", response_model=TokenResponse, operation_id="register")
async def register(
    request: Request, body: RegisterBody, conn: aiosqlite.Connection = Depends(get_db)
):
    existing = await users_db.get_by_email(conn, body.email)
    if existing:
        return error_response(request, 400, "Email already registered")
    uid = await users_db.insert_user(conn, body.email, body.display_name)
    await users_db.set_local_credentials(conn, uid, _hash_pw(body.password))
    user = await users_db.get_by_email(conn, body.email)
    assert user is not None
    access, refresh = _token_pair_for_user(user)
    resp = JSONResponse(TokenResponse(access_token=access, user=user).model_dump(by_alias=True))
    # HttpOnly refresh cookie
    resp.set_cookie(
        key="refresh_token",
        value=refresh,
        httponly=True,
        secure=False,
        samesite="lax",
        max_age=settings.refresh_ttl_seconds,
        path="/api/v1/auth/refresh",
    )
    return resp


@router.post("/login", response_model=TokenResponse, operation_id="loginV2")
async def login(request: Request, body: LoginBody, conn: aiosqlite.Connection = Depends(get_db)):
    user: Optional[User] = await users_db.get_by_email(conn, body.email)
    if not user:
        return error_response(request, 401, "Invalid credentials")
    stored = await users_db.get_local_password_hash(conn, user.id)
    if not stored or not _verify_pw(body.password, stored):
        return error_response(request, 401, "Invalid credentials")
    access, refresh = _token_pair_for_user(user)
    resp = JSONResponse(TokenResponse(access_token=access, user=user).model_dump(by_alias=True))
    resp.set_cookie(
        key="refresh_token",
        value=refresh,
        httponly=True,
        secure=False,
        samesite="lax",
        max_age=settings.refresh_ttl_seconds,
        path="/api/v1/auth/refresh",
    )
    return resp


class RefreshResponse(ApiModel):
    access_token: str
    token_type: str = "bearer"


@router.post("/refresh", response_model=RefreshResponse, operation_id="refreshV2")
async def refresh(request: Request):
    token = request.cookies.get("refresh_token")
    if not token:
        return error_response(request, 401, "Unauthorized")
    from security import verify_jwt

    try:
        _h, payload = verify_jwt(_jwt_config(), token, expected_kind="refresh")
    except Exception:
        return error_response(request, 401, "Unauthorized")
    user_id = str(payload.get("sub"))
    access = create_jwt(_jwt_config(), user_id, kind="access")
    return RefreshResponse(access_token=access)


@router.post("/logout", operation_id="logoutV2")
async def logout():
    resp = JSONResponse({"ok": True})
    resp.delete_cookie("refresh_token", path="/api/v1/auth/refresh")
    return resp