document-management/backend/dms/security/jwt_tokens.py

"""JWT 签发与严格校验。"""

from __future__ import annotations

from datetime import datetime, timedelta, timezone
from uuid import uuid4

import jwt
from flask import current_app

from dms.common.errors import AuthenticationError, InternalError, TokenExpiredError
from dms.models import User


ALGORITHM = "HS256"
REQUIRED_CLAIMS = ("sub", "username", "roleCode", "authVersion", "iat", "exp", "jti", "iss")


def _secret() -> str:
    secret = str(current_app.config.get("DMS_JWT_SECRET", "")).strip()
    if not secret:
        raise InternalError("JWT签名密钥未配置")
    return secret


def issue_access_token(user: User, *, keep_signed_in: bool) -> tuple[str, int]:
    now = datetime.now(timezone.utc)
    duration_key = (
        "DMS_JWT_KEEP_SIGNED_IN_EXPIRES_SECONDS"
        if keep_signed_in
        else "DMS_JWT_EXPIRES_SECONDS"
    )
    expires_in = int(current_app.config[duration_key])
    claims = {
        "sub": str(user.id),
        "username": user.username,
        "roleCode": user.role_code,
        "authVersion": user.auth_version,
        "iat": now,
        "exp": now + timedelta(seconds=expires_in),
        "jti": str(uuid4()),
        "iss": current_app.config["DMS_JWT_ISSUER"],
    }
    return jwt.encode(claims, _secret(), algorithm=ALGORITHM), expires_in


def decode_access_token(token: str) -> dict[str, object]:
    try:
        return jwt.decode(
            token,
            _secret(),
            algorithms=[ALGORITHM],
            issuer=current_app.config["DMS_JWT_ISSUER"],
            options={"require": list(REQUIRED_CLAIMS)},
        )
    except jwt.ExpiredSignatureError as exc:
        raise TokenExpiredError() from exc
    except jwt.PyJWTError as exc:
        raise AuthenticationError() from exc


REAUTH_TOKEN_TYPE = "reauth"
REAUTH_EXPIRES_SECONDS = 300
REAUTH_REQUIRED_CLAIMS = ("sub", "type", "authVersion", "iat", "exp", "jti", "iss")


def issue_reauth_token(user: User) -> tuple[str, int]:
    """敏感操作二次认证令牌：5 分钟有效，绑定 userId 与 authVersion（契约1.21）。"""
    now = datetime.now(timezone.utc)
    claims = {
        "sub": str(user.id),
        "type": REAUTH_TOKEN_TYPE,
        "authVersion": user.auth_version,
        "iat": now,
        "exp": now + timedelta(seconds=REAUTH_EXPIRES_SECONDS),
        "jti": str(uuid4()),
        "iss": current_app.config["DMS_JWT_ISSUER"],
    }
    return (
        jwt.encode(claims, _secret(), algorithm=ALGORITHM),
        REAUTH_EXPIRES_SECONDS,
    )


def decode_reauth_token(token: str) -> dict[str, object]:
    try:
        return jwt.decode(
            token,
            _secret(),
            algorithms=[ALGORITHM],
            issuer=current_app.config["DMS_JWT_ISSUER"],
            options={"require": list(REAUTH_REQUIRED_CLAIMS)},
        )
    except jwt.ExpiredSignatureError as exc:
        raise TokenExpiredError() from exc
    except jwt.PyJWTError as exc:
        raise AuthenticationError() from exc