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
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  # similar to Django default; adjust in settings if needed
SALT_BYTES = 16


def _hash_pw(pw: str) -> str:
    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:
    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)
        # constant-time compare
        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