meijiaka-zy/python-api/app/services/auth_service.py

"""
认证服务层
==========

封装登录、刷新、登出等核心业务逻辑。
包含 SSE 连接池管理（用于单设备登录的实时踢人）。
"""

from __future__ import annotations

import asyncio
import contextlib
import hashlib
import logging
import random
from typing import Any

import httpx
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.redis_client import get_redis_client
from app.core.security import (
    create_access_token,
    create_refresh_token,
    verify_refresh_token,
)
from app.crud.user import user as user_crud
from app.crud.user_device import user_device as device_crud
from app.models.user_device import UserDevice

logger = logging.getLogger(__name__)

# ── 短信业务常量（数值类配置不走 .env，内嵌代码）─────
SMS_CODE_LENGTH = 6              # 验证码位数
SMS_CODE_EXPIRE_MINUTES = 5      # 验证码有效期（分钟）
SMS_DAILY_LIMIT = 10             # 单手机号每日发送上限

# ========== SSE 连接池 ==========
# key: user_id, value: asyncio.Queue（用于向该用户的 SSE 连接发送消息）
_sse_connections: dict[str, asyncio.Queue[str]] = {}


def register_sse_connection(user_id: str, queue: asyncio.Queue[str]) -> None:
    """注册用户的 SSE 连接"""
    _sse_connections[user_id] = queue


def unregister_sse_connection(user_id: str) -> None:
    """注销用户的 SSE 连接"""
    _sse_connections.pop(user_id, None)


async def _kick_old_device(user_id: str) -> None:
    """
    向旧设备的 SSE 连接发送踢人消息。

    如果旧设备不在线（没有 SSE 连接），则静默忽略，
    旧设备下次请求时会发现 Token 已失效。
    """
    queue = _sse_connections.get(user_id)
    if queue is None:
        return

    with contextlib.suppress(asyncio.QueueFull):
        # 发送 kick 消息，非阻塞（旧设备可能已断开）
        queue.put_nowait('{"type": "kick", "reason": "new_device_login"}')

    # 延迟注销，给旧设备一点时间处理消息
    await asyncio.sleep(1)
    unregister_sse_connection(user_id)


# ========== 验证码校验 ==========

async def verify_sms_code(mobile: str, code: str) -> bool:
    """
    校验短信验证码。

    从 Redis 读取 sms:{mobile} 的值，与传入的 code 比对。
    校验成功后删除 Redis 中的验证码（防止重放攻击）。
    """
    redis = get_redis_client()
    key = f"sms:{mobile}"
    stored_code = await redis.get(key)

    if stored_code is None:
        return False

    if stored_code != code:
        return False

    # 校验成功，删除验证码（一次性使用）
    await redis.delete(key)
    return True


async def send_sms_code(mobile: str) -> str:
    """
    发送短信验证码。

    流程：
    1. 检查每日发送上限（Redis）
    2. 生成随机验证码
    3. 存入 Redis（带过期时间）
    4. 调用 B2M 短信平台发送（配置不完整时降级为日志打印）

    返回生成的验证码。
    """
    redis = get_redis_client()
    key = f"sms:{mobile}"

    # 检查每日发送上限
    daily_key = f"sms_daily:{mobile}"
    daily_count = await redis.get(daily_key)
    if daily_count and int(daily_count) >= SMS_DAILY_LIMIT:
        raise ValueError(f"该手机号今日验证码发送次数已达上限（{SMS_DAILY_LIMIT}次）")

    # 生成验证码
    code = "".join([str(random.randint(0, 9)) for _ in range(SMS_CODE_LENGTH)])  # nosec: B311

    # 存入 Redis，设置过期时间
    await redis.setex(key, SMS_CODE_EXPIRE_MINUTES * 60, code)

    # 更新每日计数
    ttl = await redis.ttl(daily_key)
    if ttl < 0:
        # 新key或已过期，设置24小时过期
        await redis.setex(daily_key, 24 * 60 * 60, 1)
    else:
        await redis.incr(daily_key)

    # 调用 B2M 短信平台（配置不完整时降级为日志打印）
    from app.services.sms_service import SMSError, get_sms_service

    try:
        sms = get_sms_service()
        if sms.app_id and sms.secret_key and sms.base_url:
            async with httpx.AsyncClient(timeout=httpx.Timeout(30.0, connect=10.0)) as client:
                await sms.send_verification_code(
                    client,
                    mobile=mobile,
                    code=code,
                    expire_minutes=SMS_CODE_EXPIRE_MINUTES,
                )
        else:
            # 配置不完整，记录警告但不打印验证码
            logger.warning(
                f"[SMS] B2M 短信配置不完整，验证码未发送: 手机号={mobile}"
            )
    except SMSError as e:
        logger.error(f"[SMS] 短信发送失败: {e}")
        # 短信发送失败不影响验证码生成
    except Exception as e:
        logger.exception(f"[SMS] 短信发送异常: {e}")

    return code


# ========== Token 工具 ==========

def _hash_refresh_token(token: str) -> str:
    """Refresh Token SHA256 哈希（用于数据库存储）"""
    return hashlib.sha256(token.encode()).hexdigest()


# ========== 登录服务 ==========

async def login_with_sms(
    db: AsyncSession,
    *,
    mobile: str,
    code: str,
    device_id: str,
    device_name: str | None = None,
    os_info: str | None = None,
    app_version: str | None = None,
    ip: str | None = None,
    source: str = "mobile_sms",
) -> dict[str, Any]:
    """
    手机号验证码登录。

    流程：
    1. 校验验证码
    2. 获取或创建用户
    3. 更新登录信息
    4. 踢掉旧设备（SSE 推送）
    5. 创建/覆盖设备记录
    6. 签发双 Token
    """
    # 1. 校验验证码
    if not await verify_sms_code(mobile, code):
        raise ValueError("验证码错误或已过期")

    # 2. 查询用户（不再自动注册）
    user = await user_crud.get_by_mobile(db, mobile=mobile)
    if user is None:
        raise ValueError("用户不存在")

    # 检查用户状态
    if not user.is_active:
        raise ValueError("账号已被封禁，请联系客服")

    # 3. 更新登录信息
    await user_crud.update_login_info(db, user_id=user.id, ip=ip)

    # 4. 踢掉旧设备（SSE 推送）
    await _kick_old_device(str(user.id))

    # 5. 签发双 Token
    access_token = create_access_token(data={"sub": str(user.id)})
    refresh_token = create_refresh_token(data={"sub": str(user.id)})
    refresh_token_hash = _hash_refresh_token(refresh_token)

    # 6. 创建/覆盖设备记录
    await device_crud.create_or_update(
        db,
        user_id=user.id,
        device_id=device_id,
        device_name=device_name,
        os_info=os_info,
        app_version=app_version,
        refresh_token_hash=refresh_token_hash,
    )

    return {
        "access_token": access_token,
        "refresh_token": refresh_token,
        "user": {
            "id": user.id,
            "mobile": user.mobile,
            "nickname": user.nickname or "",
            "avatar": user.avatar_url or "",
        },
    }


async def refresh_access_token(
    db: AsyncSession,
    *,
    refresh_token: str,
) -> dict[str, str]:
    """
    用 Refresh Token 换取新的 Token 对（Token 轮换）。

    流程：
    1. 校验 Refresh Token
    2. 查设备记录，匹配 refresh_token_hash
    3. 生成新 Token 对
    4. 更新设备记录的 refresh_token_hash
    """
    # 1. 校验 Refresh Token
    payload = verify_refresh_token(refresh_token)
    if not payload or not payload.get("sub"):
        raise ValueError("Refresh Token 无效或已过期")

    user_id = payload["sub"]
    refresh_token_hash = _hash_refresh_token(refresh_token)

    # 2. 查设备记录
    device = await device_crud.get_by_refresh_token_hash(
        db, refresh_token_hash=refresh_token_hash
    )
    if device is None:
        raise ValueError("设备已失效，请重新登录")

    # 3. 生成新 Token 对
    new_access_token = create_access_token(data={"sub": str(user_id)})
    new_refresh_token = create_refresh_token(data={"sub": str(user_id)})
    new_refresh_token_hash = _hash_refresh_token(new_refresh_token)

    # 4. 更新设备记录
    await device_crud.create_or_update(
        db,
        user_id=user_id,
        device_id=device.device_id,
        device_name=device.device_name,
        os_info=device.os_info,
        app_version=device.app_version,
        refresh_token_hash=new_refresh_token_hash,
    )

    return {
        "access_token": new_access_token,
        "refresh_token": new_refresh_token,
    }


async def logout(db: AsyncSession, *, user_id: str) -> bool:
    """
    用户登出。

    流程：
    1. 删除设备记录
    2. 注销 SSE 连接
    """
    await device_crud.delete_by_user_id(db, user_id=user_id)
    unregister_sse_connection(user_id)
    return True


async def get_current_user_device(
    db: AsyncSession,
    *,
    user_id: str,
) -> UserDevice | None:
    """获取当前用户的设备记录"""
    return await device_crud.get_by_user_id(db, user_id=user_id)