小鱼开发
2cece72abe
- Settings 新增 SMS_CODE_WHITELIST 配置（逗号分隔手机号） - login_with_sms 中白名单手机号跳过验证码校验 - 方便内部测试和演示账号使用
311 lines
9.0 KiB
Python
311 lines
9.0 KiB
Python
"""
|
|
认证服务层
|
|
==========
|
|
|
|
封装登录、刷新、登出等核心业务逻辑。
|
|
包含 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.config import get_settings
|
|
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. 校验验证码（白名单内的手机号跳过校验）
|
|
settings = get_settings()
|
|
if mobile not in settings.sms_code_whitelist_set and 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)
|