gateway/modules/auth/jwtService.py

# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
JWT Service
Centralizes local JWT creation and cookie helpers.
"""

from datetime import timedelta
from typing import Optional, Tuple
from fastapi import Response
from jose import jwt

from modules.shared.configuration import APP_CONFIG
from modules.shared.timeUtils import getUtcNow

# Config
SECRET_KEY = APP_CONFIG.get("APP_JWT_KEY_SECRET")
ALGORITHM = APP_CONFIG.get("Auth_ALGORITHM")
ACCESS_TOKEN_EXPIRE_MINUTES = int(APP_CONFIG.get("APP_TOKEN_EXPIRY"))
REFRESH_TOKEN_EXPIRE_DAYS = int(APP_CONFIG.get("APP_REFRESH_TOKEN_EXPIRY", "7"))

def _cookiePolicy() -> Tuple[bool, str, str]:
    """
    Return (useSecure, samesiteStarlette, samesiteSetCookieHeader).

    Evaluated on each Set-Cookie so policy is not frozen at module import (config refresh / load order).

    Cross-origin SPA + API: SameSite=None and Secure=True so credentialed fetch sends cookies.
    HTTP dev: Lax + Secure=False.

    APP_COOKIE_SECURE: explicit true/false (1/0, yes/no) overrides the APP_API_URL heuristic.
    """
    explicit = (APP_CONFIG.get("APP_COOKIE_SECURE") or "").strip().lower()
    if explicit in ("1", "true", "yes"):
        useSecure = True
    elif explicit in ("0", "false", "no"):
        useSecure = False
    else:
        apiUrl = (APP_CONFIG.get("APP_API_URL") or "").strip()
        useSecure = apiUrl.startswith("https://")
    samesite = "none" if useSecure else "lax"
    samesiteHeader = "None" if useSecure else "Lax"
    return useSecure, samesite, samesiteHeader


def createAccessToken(data: dict, expiresDelta: Optional[timedelta] = None) -> Tuple[str, "datetime"]:
    """Create a JWT access token and return (token, expiresAt)."""
    toEncode = data.copy()
    if "jti" not in toEncode or not toEncode.get("jti"):
        import uuid
        toEncode["jti"] = str(uuid.uuid4())

    expire = getUtcNow() + (expiresDelta if expiresDelta else timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES))
    toEncode.update({"exp": expire})
    encodedJwt = jwt.encode(toEncode, SECRET_KEY, algorithm=ALGORITHM)
    return encodedJwt, expire


def createRefreshToken(data: dict) -> Tuple[str, "datetime"]:
    """Create a JWT refresh token and return (token, expiresAt)."""
    toEncode = data.copy()
    if "jti" not in toEncode or not toEncode.get("jti"):
        import uuid
        toEncode["jti"] = str(uuid.uuid4())
    toEncode["type"] = "refresh"

    expire = getUtcNow() + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
    toEncode.update({"exp": expire})
    encodedJwt = jwt.encode(toEncode, SECRET_KEY, algorithm=ALGORITHM)
    return encodedJwt, expire


def setAccessTokenCookie(response: Response, token: str, expiresDelta: Optional[timedelta] = None) -> None:
    """Set access token as httpOnly cookie."""
    useSecure, samesite, _ = _cookiePolicy()
    maxAge = int(expiresDelta.total_seconds()) if expiresDelta else ACCESS_TOKEN_EXPIRE_MINUTES * 60
    response.set_cookie(
        key="auth_token",
        value=token,
        httponly=True,
        secure=useSecure,
        samesite=samesite,
        path="/",
        max_age=maxAge
    )


def setRefreshTokenCookie(response: Response, token: str) -> None:
    """Set refresh token as httpOnly cookie."""
    useSecure, samesite, _ = _cookiePolicy()
    response.set_cookie(
        key="refresh_token",
        value=token,
        httponly=True,
        secure=useSecure,
        samesite=samesite,
        path="/",
        max_age=REFRESH_TOKEN_EXPIRE_DAYS * 24 * 60 * 60
    )


def clearAccessTokenCookie(response: Response) -> None:
    """
    Clear access token cookie by setting it to expire immediately.
    Uses both raw header manipulation and FastAPI's delete_cookie for maximum browser compatibility.
    """
    useSecure, samesite, samesiteHeader = _cookiePolicy()
    secure_flag = "; Secure" if useSecure else ""

    # Primary method: Raw Set-Cookie header for guaranteed deletion
    response.headers.append(
        "Set-Cookie",
        f"auth_token=deleted; Path=/; Max-Age=0; Expires=Thu, 01 Jan 1970 00:00:00 GMT; HttpOnly{secure_flag}; SameSite={samesiteHeader}"
    )

    # Fallback: Also use FastAPI's built-in method (match SameSite/Secure for invalidation)
    response.delete_cookie(
        key="auth_token",
        path="/",
        secure=useSecure,
        httponly=True,
        samesite=samesite,
    )


def clearRefreshTokenCookie(response: Response) -> None:
    """
    Clear refresh token cookie by setting it to expire immediately.
    Uses both raw header manipulation and FastAPI's delete_cookie for maximum browser compatibility.
    """
    useSecure, samesite, samesiteHeader = _cookiePolicy()
    secure_flag = "; Secure" if useSecure else ""

    # Primary method: Raw Set-Cookie header for guaranteed deletion
    response.headers.append(
        "Set-Cookie",
        f"refresh_token=deleted; Path=/; Max-Age=0; Expires=Thu, 01 Jan 1970 00:00:00 GMT; HttpOnly{secure_flag}; SameSite={samesiteHeader}"
    )

    # Fallback: Also use FastAPI's built-in method (match SameSite/Secure for invalidation)
    response.delete_cookie(
        key="refresh_token",
        path="/",
        secure=useSecure,
        httponly=True,
        samesite=samesite,
    )