""" 认证与权限管理路由 """ from datetime import timedelta from fastapi import APIRouter, Depends, HTTPException, status, Query 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, APIKeyInfo, RegenerateAPIKeyResponse, UserCreate, ) from app.email_verification import verify_code, send_and_store_verification_code, check_rate_limit 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=["认证"]) 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, db: AsyncSession = Depends(get_db)): """ 用户/渠道/管理员/供应商登录 支持的角色： - user: 租户用户 - channel: 渠道管理员 - billing_admin: 计费管理员 - operations_admin: 运营管理员 - admin: 管理员 - super_admin: 超级管理员 - provider: 供应商管理员 """ # 根据角色查找用户 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) 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) 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, } } ) @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（可选，如果前端确保传入的是 refresh token） token_type = claims.get("type") # 为了向后兼容，不强制要求 type 为 refresh # 查询用户信息 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="注册成功" )