shanxia2.0/backend/app/core/auth.py

"""认证模块"""
import asyncio
from datetime import datetime, timedelta
from typing import Optional

import redis.asyncio as redis
from redis import exceptions as redis_exceptions
from fastapi import Depends, HTTPException, status, Request
from fastapi.security import OAuth2PasswordBearer
from jose import JWTError, jwt

from app.core.context import current_user_dept_id
from sqlalchemy.orm import Session

from app.core.config import settings
from app.core.database import get_db
from app.crud.user import user_crud
from app.core.utils import generate_uuid
from app.services.web_auth_roles import resolve_web_role_for_user

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/public/login")

REFRESH_TOKEN_EXPIRE_DAYS = 7
REFRESH_TOKEN_EXPIRE_SECONDS = int(timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS).total_seconds())

_redis_client: Optional[redis.Redis] = None


def _redis_session_enabled() -> bool:
    # 只要配置了 REDIS_HOST，就启用“会话校验/refreshToken 吊销”
    return bool(settings.REDIS_HOST)


def _get_redis_client() -> Optional[redis.Redis]:
    global _redis_client
    if not _redis_session_enabled():
        return None
    if _redis_client is None:
        _redis_client = redis.Redis(
            host=settings.REDIS_HOST,
            port=settings.REDIS_PORT,
            db=settings.REDIS_DB,
            password=settings.REDIS_PASSWORD or None,
            socket_timeout=settings.REDIS_TIMEOUT_SEC,
            ssl=settings.REDIS_SSL_ENABLED,
            decode_responses=True,
        )
    return _redis_client


def _session_key(session_id: str) -> str:
    return f"auth:sessions:{session_id}"


def _refresh_key(refresh_jti: str) -> str:
    return f"auth:refresh:{refresh_jti}"


async def store_session(session_id: str, user_id: str) -> None:
    """写入会话（用于强制失效）"""
    r = _get_redis_client()
    if r is None:
        return
    await r.set(_session_key(session_id), user_id, ex=REFRESH_TOKEN_EXPIRE_SECONDS)


async def revoke_session(session_id: str) -> None:
    """撤销会话（登出/强制失效）"""
    r = _get_redis_client()
    if r is None:
        return
    await r.delete(_session_key(session_id))


async def extend_session(session_id: str) -> None:
    """刷新会话 TTL（refreshToken 续期；已登录用户每次 API 校验通过后滑动续期）"""
    if not session_id:
        return
    r = _get_redis_client()
    if r is None:
        return
    try:
        await asyncio.wait_for(
            r.expire(_session_key(session_id), REFRESH_TOKEN_EXPIRE_SECONDS),
            timeout=3.0,
        )
    except (asyncio.TimeoutError, redis_exceptions.TimeoutError, redis_exceptions.ConnectionError):
        pass


async def is_session_valid(session_id: str) -> bool:
    """检查会话是否仍然有效"""
    r = _get_redis_client()
    if r is None:
        return True
    if not session_id:
        return False
    try:
        # 添加连接超时保护，Redis 不可用时降级为验证通过
        return await asyncio.wait_for(r.exists(_session_key(session_id)), timeout=3.0) > 0
    except (asyncio.TimeoutError, redis_exceptions.TimeoutError, redis_exceptions.ConnectionError):
        # Redis 连接失败时降级：只验证 JWT，不过滤会话吊销
        return True


async def store_refresh_token(refresh_jti: str, user_id: str) -> None:
    """写入 refreshToken（用于 one-time 使用/吊销）"""
    r = _get_redis_client()
    if r is None:
        return
    await r.set(_refresh_key(refresh_jti), user_id, ex=REFRESH_TOKEN_EXPIRE_SECONDS)


async def is_refresh_token_valid(refresh_jti: str) -> bool:
    r = _get_redis_client()
    if r is None:
        return True
    if not refresh_jti:
        return False
    try:
        return await asyncio.wait_for(r.exists(_refresh_key(refresh_jti)), timeout=3.0) > 0
    except (asyncio.TimeoutError, redis_exceptions.TimeoutError, redis_exceptions.ConnectionError):
        return True


async def revoke_refresh_token(refresh_jti: str) -> None:
    r = _get_redis_client()
    if r is None:
        return
    await r.delete(_refresh_key(refresh_jti))


def create_access_token(
    data: dict,
    session_id: str,
    token_jti: str,
    expires_delta: Optional[timedelta] = None,
) -> str:
    """创建JWT访问令牌"""
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
    to_encode.update(
        {
            "exp": expire,
            "type": "access",
            "sid": session_id,
            "jti": token_jti,
        }
    )
    encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
    return encoded_jwt


def create_refresh_token(data: dict, session_id: str, refresh_jti: str) -> str:
    """创建刷新令牌"""
    to_encode = data.copy()
    expire = datetime.utcnow() + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)  # 有效期与 Redis TTL 保持一致
    to_encode.update(
        {
            "exp": expire,
            "type": "refresh",
            "sid": session_id,
            "jti": refresh_jti,
        }
    )
    encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
    return encoded_jwt


def decode_token(token: str) -> dict:
    """解码JWT令牌"""
    try:
        payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
        return payload
    except JWTError:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="无效的认证凭据",
            headers={"WWW-Authenticate": "Bearer"},
        )


async def get_current_user_by_token(
    token: str,
    db: Session,
) -> dict:
    """根据 token 获取当前用户（包含 Redis 会话校验；会话滑动续期）"""
    payload = decode_token(token)

    user_id: str = payload.get("sub")
    session_id: str = payload.get("sid")

    if not await is_session_valid(session_id):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Token无效或已过期",
        )

    # 方案 C：有操作时延长 Redis 会话，连续 7 天无请求则会话过期
    await extend_session(session_id)

    if user_id is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="无效的认证凭据",
        )

    user = user_crud.get(db, user_id)
    if not user or getattr(user, "del_flag", 0) != 0:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="用户不存在或已删除",
        )

    if getattr(user, "status", "1") != "1":
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="用户已被禁用",
        )

    web_role = resolve_web_role_for_user(db, user.id)

    phone = (getattr(user, "phone_no", None) or "").strip() or None
    email = (getattr(user, "email", None) or "").strip() or None
    avatar = getattr(user, "avatar", None)

    current_user_dept_id.set(user.dept_id or "")

    return {
        "user_id": user.id,
        "username": user.user_no,
        "real_name": user.nick_name,
        "role": web_role,
        "region_id": user.dept_id,
        "phone": phone,
        "email": email,
        "avatar_url": avatar,
        "is_active": user.status == "1",
    }


async def get_current_user(
    token: str = Depends(oauth2_scheme),
    db: Session = Depends(get_db)
) -> dict:
    """获取当前用户"""
    return await get_current_user_by_token(token, db)


async def get_current_user_simple(token: str = Depends(oauth2_scheme)) -> dict:
    """简化版获取当前用户（不查数据库）"""
    payload = decode_token(token)
    user_id: str = payload.get("sub")
    if user_id is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="无效的认证凭据",
        )
    current_user_dept_id.set(payload.get("region_id") or "")
    return {
        "user_id": user_id,
        "username": payload.get("username"),
        "role": payload.get("role"),
        "region_id": payload.get("region_id"),
        "real_name": payload.get("real_name")
    }


def get_client_ip(request: Request) -> str:
    """获取客户端IP"""
    forwarded = request.headers.get("X-Forwarded-For")
    if forwarded:
        return forwarded.split(",")[0].strip()
    return request.client.host if request.client else "unknown"