forked from xiaohei/taiji-AI-PAD
1120 lines
37 KiB
Python
1120 lines
37 KiB
Python
"""
|
|
认证与权限管理路由
|
|
"""
|
|
|
|
from datetime import timedelta
|
|
from fastapi import APIRouter, Depends, HTTPException, status, Query, Request
|
|
from sqlalchemy import select, and_
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy.exc import IntegrityError
|
|
import secrets
|
|
import hashlib
|
|
from typing import Optional
|
|
|
|
from database import get_db
|
|
from models import (
|
|
User, Channel, APIKey, Balance, TenantCustomAgentQuota,
|
|
PlatformAgentQuota, TenantModelKey, ResourceAllocation, ModelProvider
|
|
)
|
|
from app.auth import (
|
|
authenticate_user,
|
|
create_access_token,
|
|
create_refresh_token,
|
|
get_password_hash,
|
|
verify_password,
|
|
require_auth,
|
|
)
|
|
from app.schemas import (
|
|
LoginRequest,
|
|
SuccessResponse,
|
|
TokenResponse,
|
|
PasswordChangeRequest,
|
|
ForgotPasswordRequest,
|
|
ResetPasswordRequest,
|
|
APIKeyInfo,
|
|
RegenerateAPIKeyResponse,
|
|
UserCreate,
|
|
)
|
|
from app.email_verification import verify_code, send_and_store_verification_code, check_rate_limit, send_password_reset_code
|
|
from app.audit import log_audit_event
|
|
from app.agent_manager_client import get_agent_manager_client, AgentManagerError
|
|
from app.litellm_client import get_litellm_client, LiteLLMClientError
|
|
from config import settings
|
|
import uuid
|
|
from datetime import datetime
|
|
from sqlalchemy import and_
|
|
import structlog
|
|
|
|
logger = structlog.get_logger(__name__)
|
|
|
|
router = APIRouter(prefix="/api/auth", tags=["认证"])
|
|
|
|
|
|
class _LoginRateLimit:
|
|
"""每 IP 维度对登录请求限流（默认 5 次/60s，超出 429）。
|
|
多副本部署下每副本独立计数，可接受。"""
|
|
def __init__(self, max_attempts: int = 5, window_seconds: int = 60):
|
|
self.max_attempts = max_attempts
|
|
self.window_seconds = window_seconds
|
|
self._attempts: dict = {}
|
|
|
|
def _client_ip(self, request) -> str:
|
|
xff = request.headers.get("x-forwarded-for")
|
|
if xff:
|
|
return xff.split(",")[0].strip()
|
|
return request.client.host if request.client else "unknown"
|
|
|
|
async def __call__(self, request: Request):
|
|
import time
|
|
from collections import deque
|
|
ip = self._client_ip(request)
|
|
now = time.time()
|
|
bucket = self._attempts.setdefault(ip, deque())
|
|
while bucket and bucket[0] < now - self.window_seconds:
|
|
bucket.popleft()
|
|
if len(bucket) >= self.max_attempts:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail="登录尝试过于频繁，请稍后再试",
|
|
headers={"Retry-After": str(self.window_seconds)},
|
|
)
|
|
bucket.append(now)
|
|
|
|
|
|
_login_rate_limit = _LoginRateLimit()
|
|
|
|
|
|
def _mask_api_key(key: str) -> str:
|
|
"""隐藏API密钥的中间部分"""
|
|
if len(key) <= 12:
|
|
return key[:4] + "..." + key[-4:]
|
|
return key[:8] + "..." + key[-4:]
|
|
|
|
|
|
@router.post("/login", response_model=SuccessResponse)
|
|
async def login(req: LoginRequest, request: Request, db: AsyncSession = Depends(get_db), _: None = Depends(_login_rate_limit)):
|
|
"""
|
|
用户/渠道/管理员/供应商登录
|
|
|
|
支持的角色：
|
|
- user: 租户用户
|
|
- channel: 渠道管理员
|
|
- billing_admin: 计费管理员
|
|
- operations_admin: 运营管理员
|
|
- admin: 管理员
|
|
- super_admin: 超级管理员
|
|
- provider: 供应商管理员
|
|
"""
|
|
success = False
|
|
result_user_id: Optional[str] = None
|
|
error_msg: Optional[str] = None
|
|
try:
|
|
# 根据角色查找用户
|
|
if req.role == "channel":
|
|
# 渠道管理员登录
|
|
result = await db.execute(select(Channel).where(Channel.email == req.email))
|
|
entity = result.scalar_one_or_none()
|
|
|
|
if not entity or not verify_password(req.password, entity.password_hash):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="邮箱或密码错误"
|
|
)
|
|
|
|
# 更新最后登录时间
|
|
from datetime import datetime
|
|
entity.last_login_at = datetime.utcnow()
|
|
await db.commit()
|
|
|
|
# 创建JWT token
|
|
token_data = {
|
|
"sub": str(entity.id),
|
|
"email": entity.email,
|
|
"role": "channel_admin",
|
|
"channelId": str(entity.id),
|
|
}
|
|
access_token = create_access_token(data=token_data)
|
|
refresh_token = create_refresh_token(data=token_data)
|
|
|
|
success = True
|
|
result_user_id = str(entity.id)
|
|
return SuccessResponse(
|
|
data={
|
|
"token": access_token,
|
|
"refreshToken": refresh_token,
|
|
"user": {
|
|
"id": str(entity.id),
|
|
"name": entity.name,
|
|
"email": entity.email,
|
|
"role": "channel_admin",
|
|
"channelId": str(entity.id),
|
|
}
|
|
}
|
|
)
|
|
|
|
else:
|
|
# 用户/管理员/供应商登录
|
|
result = await db.execute(select(User).where(User.email == req.email))
|
|
user = result.scalar_one_or_none()
|
|
|
|
if not user:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="邮箱或密码错误"
|
|
)
|
|
|
|
# 验证密码（兼容两种密码字段）
|
|
password_hash = user.password_hash or user.hashed_password
|
|
if not password_hash or not verify_password(req.password, password_hash):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="邮箱或密码错误"
|
|
)
|
|
|
|
# 验证用户状态
|
|
if hasattr(user, 'status') and user.status == "inactive":
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="账户已被禁用"
|
|
)
|
|
|
|
# 验证角色
|
|
user_role = user.role
|
|
|
|
# 角色验证逻辑
|
|
valid_roles = {
|
|
"super_admin": ["super_admin"],
|
|
"admin": ["admin", "super_admin"],
|
|
"billing_admin": ["billing_admin", "admin", "super_admin"],
|
|
"operations_admin": ["operations_admin", "admin", "super_admin"],
|
|
"user": ["user"],
|
|
"provider": ["provider_admin"],
|
|
}
|
|
|
|
allowed_roles = valid_roles.get(req.role, [])
|
|
if user_role not in allowed_roles:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail=f"权限不足，当前角色: {user_role}"
|
|
)
|
|
|
|
# 更新最后登录时间
|
|
from datetime import datetime
|
|
user.last_login_at = datetime.utcnow()
|
|
await db.commit()
|
|
|
|
# 创建JWT token
|
|
token_data = {
|
|
"sub": str(user.id),
|
|
"email": user.email,
|
|
"role": user_role,
|
|
"channelId": str(user.channel_id) if user.channel_id else None,
|
|
}
|
|
access_token = create_access_token(data=token_data)
|
|
refresh_token = create_refresh_token(data=token_data)
|
|
|
|
success = True
|
|
result_user_id = str(user.id)
|
|
return SuccessResponse(
|
|
data={
|
|
"token": access_token,
|
|
"refreshToken": refresh_token,
|
|
"user": {
|
|
"id": str(user.id),
|
|
"name": user.name or user.full_name,
|
|
"email": user.email,
|
|
"role": user_role,
|
|
"channelId": str(user.channel_id) if user.channel_id else None,
|
|
}
|
|
}
|
|
)
|
|
except HTTPException as e:
|
|
error_msg = e.detail if isinstance(e.detail, str) else str(e.detail)
|
|
raise
|
|
finally:
|
|
try:
|
|
await log_audit_event(
|
|
action="auth.login",
|
|
resource_type="user",
|
|
resource_id=req.email,
|
|
user_id=result_user_id,
|
|
success=success,
|
|
details={"role": req.role},
|
|
error_message=error_msg,
|
|
request=request,
|
|
db=db,
|
|
)
|
|
except Exception as audit_exc:
|
|
logger.warning("auth_login_audit_failed", error=str(audit_exc))
|
|
|
|
|
|
@router.get("/me", response_model=SuccessResponse)
|
|
async def get_current_user(
|
|
principal: dict = Depends(require_auth),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
"""
|
|
获取当前登录用户信息
|
|
"""
|
|
user_id = principal.get("user_id")
|
|
result = await db.execute(select(User).where(User.id == user_id))
|
|
user = result.scalar_one_or_none()
|
|
|
|
if user is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="用户不存在"
|
|
)
|
|
|
|
if getattr(user, "status", None) == "inactive":
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="账户已被禁用"
|
|
)
|
|
|
|
return SuccessResponse(
|
|
data={
|
|
"id": str(user.id),
|
|
"email": user.email,
|
|
"name": user.name or user.full_name,
|
|
"role": user.role,
|
|
"channelId": str(user.channel_id) if user.channel_id is not None else None,
|
|
"status": user.status,
|
|
"subscriptionTier": getattr(user, "subscription_tier", None),
|
|
"lastLoginAt": user.last_login_at.isoformat() if user.last_login_at is not None else None,
|
|
}
|
|
)
|
|
|
|
|
|
@router.post("/logout", response_model=SuccessResponse)
|
|
async def logout(
|
|
principal: dict = Depends(require_auth),
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
"""
|
|
用户登出
|
|
|
|
将当前Token加入黑名单
|
|
"""
|
|
from datetime import datetime, timedelta
|
|
from app.token_blacklist import add_token_to_blacklist
|
|
import uuid as uuid_module
|
|
|
|
user_id = principal.get("user_id")
|
|
claims = principal.get("claims", {})
|
|
|
|
# 确保 user_id 是有效的 UUID 字符串
|
|
if user_id:
|
|
try:
|
|
# 验证并转换为标准 UUID 字符串格式
|
|
user_id = str(uuid_module.UUID(str(user_id)))
|
|
except (ValueError, TypeError):
|
|
# 如果无法转换，使用 claims 中的 sub
|
|
sub = claims.get("sub")
|
|
if sub:
|
|
try:
|
|
user_id = str(uuid_module.UUID(str(sub)))
|
|
except (ValueError, TypeError):
|
|
user_id = None
|
|
|
|
if not user_id:
|
|
# 如果无法获取有效的 user_id，直接返回成功（前端会清除本地 token）
|
|
return SuccessResponse(message="登出成功")
|
|
|
|
# 获取Token的过期时间（从claims中提取）
|
|
exp = claims.get("exp")
|
|
if exp:
|
|
expires_at = datetime.utcfromtimestamp(exp)
|
|
else:
|
|
# 默认7天后过期
|
|
expires_at = datetime.utcnow() + timedelta(days=7)
|
|
|
|
# 生成token的唯一标识（使用sub + 时间戳）
|
|
token_jti = f"logout_{user_id}_{datetime.utcnow().timestamp()}"
|
|
|
|
try:
|
|
# 添加到黑名单
|
|
await add_token_to_blacklist(
|
|
token_jti=token_jti,
|
|
user_id=user_id,
|
|
expires_at=expires_at,
|
|
reason="logout",
|
|
db=db,
|
|
)
|
|
except Exception as e:
|
|
# 即使黑名单添加失败，也返回成功（前端会清除本地 token）
|
|
import structlog
|
|
logger = structlog.get_logger(__name__)
|
|
logger.warning("failed_to_add_token_to_blacklist", error=str(e), user_id=user_id)
|
|
|
|
return SuccessResponse(message="登出成功")
|
|
|
|
|
|
@router.post("/refresh", response_model=SuccessResponse)
|
|
async def refresh_token_endpoint(
|
|
principal: dict = Depends(require_auth),
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
"""
|
|
刷新访问令牌
|
|
|
|
使用 Refresh Token 获取新的 Access Token
|
|
"""
|
|
user_id = principal.get("user_id")
|
|
claims = principal.get("claims", {})
|
|
|
|
# 验证是否为 refresh token
|
|
if claims.get("type") != "refresh":
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="必须使用 refresh token 调用此接口"
|
|
)
|
|
|
|
# 查询用户信息
|
|
result = await db.execute(select(User).where(User.id == user_id))
|
|
user = result.scalar_one_or_none()
|
|
|
|
if not user:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="用户不存在"
|
|
)
|
|
|
|
# 生成新的 access token 和 refresh token
|
|
token_data = {
|
|
"sub": str(user.id),
|
|
"email": user.email,
|
|
"role": user.role,
|
|
"channelId": str(user.channel_id) if user.channel_id else None,
|
|
}
|
|
new_access_token = create_access_token(data=token_data)
|
|
new_refresh_token = create_refresh_token(data=token_data)
|
|
|
|
return SuccessResponse(
|
|
data={
|
|
"token": new_access_token,
|
|
"refreshToken": new_refresh_token,
|
|
}
|
|
)
|
|
|
|
|
|
@router.put("/password", response_model=SuccessResponse)
|
|
async def change_password(
|
|
req: PasswordChangeRequest,
|
|
principal: dict = Depends(require_auth),
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
"""
|
|
修改密码
|
|
"""
|
|
user_id = principal.get("user_id")
|
|
|
|
# 查询用户
|
|
result = await db.execute(select(User).where(User.id == user_id))
|
|
user = result.scalar_one_or_none()
|
|
|
|
if not user:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="用户不存在"
|
|
)
|
|
|
|
# 验证旧密码
|
|
password_hash = user.password_hash or user.hashed_password
|
|
if not verify_password(req.old_password, password_hash):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="旧密码错误"
|
|
)
|
|
|
|
# 更新密码
|
|
new_hash = get_password_hash(req.new_password)
|
|
user.password_hash = new_hash
|
|
user.hashed_password = new_hash # 兼容旧字段
|
|
|
|
await db.commit()
|
|
|
|
return SuccessResponse(message="密码修改成功")
|
|
|
|
|
|
@router.get("/keys/info", response_model=SuccessResponse)
|
|
async def get_api_key_info(
|
|
principal: dict = Depends(require_auth),
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
"""
|
|
获取服务终结点和API密钥信息
|
|
"""
|
|
user_id = principal.get("user_id")
|
|
|
|
# 查询用户的API密钥
|
|
result = await db.execute(
|
|
select(APIKey)
|
|
.where(APIKey.user_id == user_id)
|
|
.where(APIKey.is_active == True)
|
|
.order_by(APIKey.created_at.desc())
|
|
)
|
|
api_key = result.scalar_one_or_none()
|
|
|
|
if not api_key:
|
|
# 创建新的API密钥
|
|
raw_key = f"sk-{secrets.token_urlsafe(32)}"
|
|
key_hash = get_password_hash(raw_key)
|
|
|
|
new_key = APIKey(
|
|
user_id=user_id,
|
|
api_key_hash=key_hash,
|
|
api_key_prefix=raw_key[:8],
|
|
name="默认密钥",
|
|
key_hash=key_hash,
|
|
prefix=raw_key[:8],
|
|
is_active=True,
|
|
)
|
|
db.add(new_key)
|
|
await db.commit()
|
|
await db.refresh(new_key)
|
|
|
|
api_key = new_key
|
|
masked_key = _mask_api_key(raw_key)
|
|
else:
|
|
# 隐藏现有密钥
|
|
masked_key = api_key.api_key_prefix + "..." + "xxxx"
|
|
|
|
endpoint = f"https://api.taiji-ai.com/v1" # 可以从配置读取
|
|
|
|
return SuccessResponse(
|
|
data={
|
|
"endpoint": endpoint,
|
|
"apiKey": masked_key,
|
|
"createdAt": api_key.created_at.isoformat(),
|
|
"lastUsed": api_key.last_used.isoformat() if api_key.last_used else None,
|
|
}
|
|
)
|
|
|
|
|
|
@router.post("/keys/regenerate", response_model=SuccessResponse)
|
|
async def regenerate_api_key(
|
|
principal: dict = Depends(require_auth),
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
"""
|
|
重新生成API密钥
|
|
"""
|
|
user_id = principal.get("user_id")
|
|
|
|
# 禁用所有旧密钥
|
|
result = await db.execute(
|
|
select(APIKey).where(APIKey.user_id == user_id)
|
|
)
|
|
old_keys = result.scalars().all()
|
|
|
|
for key in old_keys:
|
|
key.is_active = False
|
|
|
|
# 生成新密钥
|
|
raw_key = f"sk-{secrets.token_urlsafe(32)}"
|
|
key_hash = get_password_hash(raw_key)
|
|
|
|
new_key = APIKey(
|
|
user_id=user_id,
|
|
api_key_hash=key_hash,
|
|
api_key_prefix=raw_key[:8],
|
|
name="重新生成的密钥",
|
|
key_hash=key_hash,
|
|
prefix=raw_key[:8],
|
|
is_active=True,
|
|
)
|
|
db.add(new_key)
|
|
await db.commit()
|
|
|
|
return SuccessResponse(
|
|
data={
|
|
"apiKey": raw_key, # 完整显示新密钥
|
|
"message": "旧密钥已失效",
|
|
}
|
|
)
|
|
|
|
|
|
@router.get("/keys", response_model=SuccessResponse)
|
|
async def list_api_keys(
|
|
principal: dict = Depends(require_auth),
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
"""
|
|
获取用户的 API 密钥列表
|
|
|
|
返回所有密钥的信息（不包含完整密钥内容，仅显示前缀）
|
|
"""
|
|
user_id = principal.get("user_id")
|
|
|
|
result = await db.execute(
|
|
select(APIKey)
|
|
.where(APIKey.user_id == user_id)
|
|
.order_by(APIKey.created_at.desc())
|
|
)
|
|
keys = result.scalars().all()
|
|
|
|
return SuccessResponse(
|
|
data={
|
|
"keys": [
|
|
{
|
|
"id": str(key.id),
|
|
"name": key.name or "默认密钥",
|
|
"prefix": key.api_key_prefix + "...",
|
|
"isActive": key.is_active,
|
|
"createdAt": key.created_at.isoformat(),
|
|
"lastUsed": key.last_used.isoformat() if key.last_used else None,
|
|
"expiresAt": key.expires_at.isoformat() if key.expires_at else None,
|
|
"totalRequests": key.total_requests or 0,
|
|
}
|
|
for key in keys
|
|
],
|
|
"total": len(keys)
|
|
},
|
|
message=f"共 {len(keys)} 个 API 密钥"
|
|
)
|
|
|
|
|
|
@router.post("/keys", response_model=SuccessResponse)
|
|
async def create_api_key(
|
|
name: str = Query(default="API Key", description="密钥名称"),
|
|
expires_in_days: Optional[int] = Query(default=None, description="过期天数，不填则永不过期"),
|
|
principal: dict = Depends(require_auth),
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
"""
|
|
创建新的 API 密钥
|
|
|
|
注意：密钥只在创建时显示一次，请妥善保管
|
|
|
|
参数：
|
|
- name: 密钥名称，用于标识不同用途的密钥
|
|
- expires_in_days: 过期天数，不填则永不过期
|
|
|
|
返回：
|
|
- key: 完整的 API 密钥（只在创建时返回一次）
|
|
- id: 密钥 ID，用于删除操作
|
|
"""
|
|
user_id = principal.get("user_id")
|
|
|
|
# 生成新密钥
|
|
raw_key = f"sk-{secrets.token_urlsafe(32)}"
|
|
key_hash = get_password_hash(raw_key)
|
|
|
|
expires_at = None
|
|
if expires_in_days:
|
|
expires_at = datetime.utcnow() + timedelta(days=expires_in_days)
|
|
|
|
new_key = APIKey(
|
|
user_id=user_id,
|
|
api_key_hash=key_hash,
|
|
api_key_prefix=raw_key[:8],
|
|
name=name,
|
|
key_hash=key_hash,
|
|
prefix=raw_key[:8],
|
|
is_active=True,
|
|
expires_at=expires_at,
|
|
total_requests=0,
|
|
)
|
|
db.add(new_key)
|
|
await db.commit()
|
|
await db.refresh(new_key)
|
|
|
|
return SuccessResponse(
|
|
data={
|
|
"id": str(new_key.id),
|
|
"name": name,
|
|
"key": raw_key, # 只在创建时返回完整密钥
|
|
"prefix": raw_key[:8] + "...",
|
|
"expiresAt": expires_at.isoformat() if expires_at else None,
|
|
"createdAt": new_key.created_at.isoformat(),
|
|
},
|
|
message="API 密钥创建成功，请妥善保管，密钥只显示一次"
|
|
)
|
|
|
|
|
|
@router.delete("/keys/{key_id}", response_model=SuccessResponse)
|
|
async def delete_api_key(
|
|
key_id: str,
|
|
principal: dict = Depends(require_auth),
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
"""
|
|
删除 API 密钥
|
|
|
|
参数：
|
|
- key_id: 密钥 ID（从列表接口或创建接口获取）
|
|
|
|
注意：删除后密钥立即失效，无法恢复
|
|
"""
|
|
user_id = principal.get("user_id")
|
|
|
|
try:
|
|
key_uuid = uuid.UUID(key_id)
|
|
except ValueError:
|
|
raise HTTPException(status_code=400, detail="无效的密钥 ID")
|
|
|
|
result = await db.execute(
|
|
select(APIKey)
|
|
.where(APIKey.id == key_uuid)
|
|
.where(APIKey.user_id == user_id)
|
|
)
|
|
key = result.scalar_one_or_none()
|
|
|
|
if not key:
|
|
raise HTTPException(status_code=404, detail="密钥不存在或无权删除")
|
|
|
|
await db.delete(key)
|
|
await db.commit()
|
|
|
|
return SuccessResponse(
|
|
data={"id": key_id, "name": key.name},
|
|
message=f"API 密钥 '{key.name}' 已删除"
|
|
)
|
|
|
|
|
|
@router.post("/register/send-code", response_model=SuccessResponse)
|
|
async def send_verification_code_endpoint(
|
|
email: str = Query(..., description="邮箱地址"),
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
"""
|
|
发送邮箱验证码
|
|
|
|
在用户注册前，先调用此接口发送验证码到邮箱
|
|
|
|
频率限制：同一邮箱60秒内只能发送一次
|
|
"""
|
|
# 1. 检查发送频率限制
|
|
can_send, remaining_seconds = await check_rate_limit(email)
|
|
if not can_send:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail=f"请等待{remaining_seconds}秒后再重新发送验证码"
|
|
)
|
|
|
|
# 2. 检查邮箱是否已存在
|
|
result = await db.execute(select(User).where(User.email == email))
|
|
existing_user = result.scalar_one_or_none()
|
|
|
|
if existing_user:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="该邮箱已被注册"
|
|
)
|
|
|
|
# 3. 发送验证码
|
|
code = await send_and_store_verification_code(email)
|
|
if not code:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="验证码发送失败，请稍后重试"
|
|
)
|
|
|
|
return SuccessResponse(
|
|
message="验证码已发送到您的邮箱，请查收"
|
|
)
|
|
|
|
|
|
@router.post("/register", response_model=SuccessResponse)
|
|
async def register(req: UserCreate, db: AsyncSession = Depends(get_db)):
|
|
"""
|
|
用户注册接口
|
|
|
|
允许用户自由注册，创建普通用户账户
|
|
注册成功后自动分配默认资源：
|
|
- 分配到 taiji 渠道 (channelId: b415e70b-8d37-481c-b229-bc3b7871607b)
|
|
- 自定义 agent 配额：2 CPU, 2 GB 内存
|
|
- 平台 agent 各1个
|
|
- 供应商所有模型
|
|
- 余额：20元
|
|
"""
|
|
# 1. 先检查邮箱是否已存在（不消耗验证码）
|
|
result = await db.execute(select(User).where(User.email == req.email))
|
|
existing_user = result.scalar_one_or_none()
|
|
|
|
if existing_user:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="该邮箱已被注册"
|
|
)
|
|
|
|
# 2. 检查用户名是否已存在（username 是必填字段）
|
|
result = await db.execute(select(User).where(User.username == req.username))
|
|
existing_username = result.scalar_one_or_none()
|
|
if existing_username:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="该用户名已被使用"
|
|
)
|
|
|
|
# 3. 最后验证邮箱验证码（验证成功后会消耗验证码）
|
|
is_valid = await verify_code(req.email, req.verification_code)
|
|
if not is_valid:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="验证码错误或已过期"
|
|
)
|
|
|
|
# taiji 渠道 ID
|
|
TAIJI_CHANNEL_ID = uuid.UUID("b415e70b-8d37-481c-b229-bc3b7871607b")
|
|
|
|
# 验证渠道存在
|
|
channel_result = await db.execute(
|
|
select(Channel).where(Channel.id == TAIJI_CHANNEL_ID)
|
|
)
|
|
channel = channel_result.scalar_one_or_none()
|
|
if not channel:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="渠道不存在，请联系管理员"
|
|
)
|
|
|
|
# 创建新用户
|
|
password_hash = get_password_hash(req.password)
|
|
username = req.username or req.email.split("@")[0]
|
|
name = req.full_name or username
|
|
|
|
new_user = User(
|
|
name=name,
|
|
email=req.email,
|
|
password_hash=password_hash,
|
|
hashed_password=password_hash, # 兼容字段
|
|
username=username,
|
|
full_name=req.full_name or name,
|
|
role="user", # 默认角色为普通用户
|
|
channel_id=TAIJI_CHANNEL_ID, # 分配到 taiji 渠道
|
|
status="active",
|
|
is_active=True,
|
|
is_admin=False,
|
|
credit_limit=0,
|
|
total_eu_consumed=0,
|
|
)
|
|
|
|
try:
|
|
db.add(new_user)
|
|
await db.flush() # 获取 user.id，这里会触发唯一约束检查
|
|
except IntegrityError as e:
|
|
await db.rollback()
|
|
error_str = str(e.orig) if e.orig else str(e)
|
|
if "email" in error_str.lower() or "users_email_key" in error_str.lower():
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="该邮箱已被注册"
|
|
)
|
|
elif "username" in error_str.lower() or "users_username_key" in error_str.lower():
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="该用户名已被使用"
|
|
)
|
|
else:
|
|
logger.error(f"注册时数据库约束冲突: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="注册信息与现有用户冲突，请更换邮箱或用户名"
|
|
)
|
|
|
|
user_id = new_user.id
|
|
|
|
try:
|
|
# 1. 创建余额记录，初始余额 20 元
|
|
balance = Balance(
|
|
user_id=user_id,
|
|
eu_balance=20.0
|
|
)
|
|
db.add(balance)
|
|
|
|
# 2. 分配自定义 Agent 配额：2 CPU, 2 GB 内存
|
|
custom_agent_quota = TenantCustomAgentQuota(
|
|
tenant_id=user_id,
|
|
cpu_quota=2.0,
|
|
memory_quota=2.0,
|
|
cpu_used=0.0,
|
|
memory_used=0.0,
|
|
agent_count=0
|
|
)
|
|
db.add(custom_agent_quota)
|
|
|
|
# 3. 获取所有平台 Agent 模板，为每个模板分配 1 个配额
|
|
try:
|
|
agent_manager_client = get_agent_manager_client()
|
|
platform_templates = await agent_manager_client.list_platform_templates()
|
|
|
|
for template in platform_templates:
|
|
template_name = template.template
|
|
|
|
# 获取渠道的配额记录（需要同步更新 pod_used）
|
|
channel_quota_result = await db.execute(
|
|
select(PlatformAgentQuota)
|
|
.where(
|
|
and_(
|
|
PlatformAgentQuota.target_id == TAIJI_CHANNEL_ID,
|
|
PlatformAgentQuota.target_type == "channel",
|
|
PlatformAgentQuota.template_name == template_name
|
|
)
|
|
)
|
|
.with_for_update() # 行锁
|
|
)
|
|
channel_quota = channel_quota_result.scalar_one_or_none()
|
|
|
|
# 如果渠道没有该模板的配额，跳过
|
|
if not channel_quota:
|
|
logger.warning(f"渠道没有平台 Agent '{template_name}' 的配额，跳过分配")
|
|
continue
|
|
|
|
# 检查渠道剩余配额是否足够
|
|
remaining = channel_quota.pod_quota - (channel_quota.pod_used or 0)
|
|
if remaining < 1:
|
|
logger.warning(f"渠道平台 Agent '{template_name}' 配额不足，跳过分配")
|
|
continue
|
|
|
|
# 创建平台 Agent 配额记录
|
|
platform_quota = PlatformAgentQuota(
|
|
target_id=user_id,
|
|
target_type="tenant",
|
|
template_name=template_name,
|
|
pod_quota=1, # 每个平台 agent 分配 1 个配额
|
|
pod_used=0,
|
|
allocated_by=None, # 系统自动分配
|
|
allocated_at=datetime.utcnow()
|
|
)
|
|
db.add(platform_quota)
|
|
|
|
# Bug 修复：更新渠道的 pod_used（分配给租户的配额视为渠道已使用的配额）
|
|
channel_quota.pod_used = (channel_quota.pod_used or 0) + 1
|
|
|
|
# 同时创建 ResourceAllocation 记录（兼容旧逻辑）
|
|
allocation = ResourceAllocation(
|
|
target_id=user_id,
|
|
target_type="tenant",
|
|
resource_type="agent",
|
|
resource_id=template_name,
|
|
quantity=1
|
|
)
|
|
db.add(allocation)
|
|
except (AgentManagerError, Exception) as e:
|
|
logger.warning(f"获取平台 Agent 模板失败，跳过平台 Agent 分配: {e}")
|
|
|
|
# 4. 分配所有供应商模型
|
|
providers_result = await db.execute(
|
|
select(ModelProvider).where(ModelProvider.is_active == True)
|
|
)
|
|
providers = providers_result.scalars().all()
|
|
|
|
if channel.litellm_team_id:
|
|
litellm_client = get_litellm_client()
|
|
|
|
for provider in providers:
|
|
# 为每个供应商的每个模型创建 TenantModelKey
|
|
for model_name in provider.supported_models:
|
|
try:
|
|
# 在 LiteLLM 中创建 Key
|
|
key = await litellm_client.generate_key(
|
|
team_id=channel.litellm_team_id,
|
|
models=[model_name],
|
|
rpm_limit=provider.rpm or 60,
|
|
tpm_limit=provider.tpm or 10000,
|
|
max_budget=500.0, # 默认预算
|
|
budget_duration="monthly",
|
|
key_name=f"tenant-{user_id}-{model_name}",
|
|
metadata={
|
|
"tenant_id": str(user_id),
|
|
"tenant_name": name,
|
|
"channel_id": str(TAIJI_CHANNEL_ID),
|
|
"channel_name": channel.name,
|
|
"model": model_name,
|
|
}
|
|
)
|
|
|
|
# 加密存储 Key
|
|
encrypted_key = litellm_client.encrypt_key(key.key)
|
|
|
|
# 保存到数据库
|
|
tenant_key = TenantModelKey(
|
|
tenant_id=user_id,
|
|
channel_id=TAIJI_CHANNEL_ID,
|
|
model_name=model_name,
|
|
litellm_key_id=key.key,
|
|
litellm_key_hash=encrypted_key,
|
|
rpm_limit=provider.rpm or 60,
|
|
tpm_limit=provider.tpm or 10000,
|
|
max_budget=500.0,
|
|
budget_duration="monthly",
|
|
status="active",
|
|
)
|
|
db.add(tenant_key)
|
|
|
|
# 同时记录到 ResourceAllocation
|
|
model_allocation = ResourceAllocation(
|
|
target_id=user_id,
|
|
target_type="tenant",
|
|
resource_type="model",
|
|
resource_id=str(provider.id),
|
|
rpm=provider.rpm or 60,
|
|
tpm=provider.tpm or 10000,
|
|
)
|
|
db.add(model_allocation)
|
|
except (LiteLLMClientError, Exception) as e:
|
|
logger.warning(f"为模型 {model_name} 创建 LiteLLM Key 失败: {e}")
|
|
# 继续处理其他模型
|
|
else:
|
|
logger.warning(f"渠道 {TAIJI_CHANNEL_ID} 未配置 LiteLLM team，跳过模型分配")
|
|
|
|
# 提交所有更改
|
|
await db.commit()
|
|
await db.refresh(new_user)
|
|
|
|
logger.info(
|
|
f"用户注册成功并分配默认资源",
|
|
user_id=str(user_id),
|
|
email=req.email,
|
|
channel_id=str(TAIJI_CHANNEL_ID)
|
|
)
|
|
|
|
except Exception as e:
|
|
await db.rollback()
|
|
logger.error(f"用户注册失败: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail=f"注册失败: {str(e)}"
|
|
)
|
|
|
|
# 创建JWT token，自动登录
|
|
token_data = {
|
|
"sub": str(new_user.id),
|
|
"email": new_user.email,
|
|
"role": new_user.role,
|
|
"user_id": str(new_user.id),
|
|
"channelId": str(TAIJI_CHANNEL_ID),
|
|
}
|
|
|
|
# 创建 Access Token（短期有效）和 Refresh Token（长期有效）
|
|
access_token = create_access_token(data=token_data)
|
|
refresh_token = create_refresh_token(data=token_data)
|
|
|
|
return SuccessResponse(
|
|
data={
|
|
"token": access_token,
|
|
"refreshToken": refresh_token,
|
|
"user": {
|
|
"id": str(new_user.id),
|
|
"name": new_user.name,
|
|
"email": new_user.email,
|
|
"username": new_user.username,
|
|
"role": new_user.role,
|
|
"channelId": str(TAIJI_CHANNEL_ID),
|
|
}
|
|
},
|
|
message="注册成功"
|
|
)
|
|
|
|
|
|
@router.post("/forgot-password/send-code", response_model=SuccessResponse)
|
|
async def forgot_password_send_code(
|
|
req: ForgotPasswordRequest,
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
"""
|
|
忘记密码 - 发送验证码
|
|
|
|
向用户注册邮箱发送密码重置验证码。
|
|
|
|
频率限制：同一邮箱60秒内只能发送一次
|
|
|
|
请求体：
|
|
{
|
|
"email": "user@example.com"
|
|
}
|
|
|
|
响应：
|
|
- 成功：返回验证码已发送的消息
|
|
- 失败：邮箱不存在、发送频率限制等
|
|
"""
|
|
# 1. 检查发送频率限制
|
|
can_send, remaining_seconds = await check_rate_limit(req.email)
|
|
if not can_send:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail=f"请等待{remaining_seconds}秒后再重新发送验证码"
|
|
)
|
|
|
|
# 2. 检查邮箱是否存在（必须是已注册用户）
|
|
result = await db.execute(select(User).where(User.email == req.email))
|
|
user = result.scalar_one_or_none()
|
|
|
|
if not user:
|
|
# 为了安全，不直接告知邮箱不存在，但记录日志
|
|
logger.warning("忘记密码请求：邮箱不存在", email=req.email)
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="该邮箱未注册"
|
|
)
|
|
|
|
# 3. 发送密码重置验证码（使用专门的密码重置邮件模板）
|
|
code = await send_password_reset_code(req.email)
|
|
if not code:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail="验证码发送失败，请稍后重试"
|
|
)
|
|
|
|
logger.info("忘记密码验证码已发送", email=req.email, user_id=str(user.id))
|
|
|
|
return SuccessResponse(
|
|
message="验证码已发送到您的邮箱，请查收"
|
|
)
|
|
|
|
|
|
@router.post("/forgot-password/reset", response_model=SuccessResponse)
|
|
async def forgot_password_reset(
|
|
req: ResetPasswordRequest,
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
"""
|
|
忘记密码 - 重置密码
|
|
|
|
验证邮箱验证码并重置密码。
|
|
|
|
请求体：
|
|
{
|
|
"email": "user@example.com",
|
|
"verification_code": "123456",
|
|
"new_password": "newPassword123"
|
|
}
|
|
|
|
响应：
|
|
- 成功：密码重置成功
|
|
- 失败：验证码错误、邮箱不存在等
|
|
"""
|
|
# 1. 检查邮箱是否存在
|
|
result = await db.execute(select(User).where(User.email == req.email))
|
|
user = result.scalar_one_or_none()
|
|
|
|
if not user:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="该邮箱未注册"
|
|
)
|
|
|
|
# 2. 验证验证码
|
|
is_valid = await verify_code(req.email, req.verification_code)
|
|
if not is_valid:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="验证码错误或已过期"
|
|
)
|
|
|
|
# 3. 更新密码
|
|
new_hash = get_password_hash(req.new_password)
|
|
user.password_hash = new_hash
|
|
user.hashed_password = new_hash # 兼容旧字段
|
|
|
|
await db.commit()
|
|
|
|
logger.info("密码重置成功", email=req.email, user_id=str(user.id))
|
|
|
|
return SuccessResponse(
|
|
message="密码重置成功，请使用新密码登录"
|
|
)
|
|
|