package controller import ( "errors" "fmt" "net/http" "one-api/common" "one-api/model" "strconv" "github.com/gin-contrib/sessions" "github.com/gin-gonic/gin" ) // Setup2FARequest 设置2FA请求结构 type Setup2FARequest struct { Code string `json:"code" binding:"required"` } // Verify2FARequest 验证2FA请求结构 type Verify2FARequest struct { Code string `json:"code" binding:"required"` } // Setup2FAResponse 设置2FA响应结构 type Setup2FAResponse struct { Secret string `json:"secret"` QRCodeData string `json:"qr_code_data"` BackupCodes []string `json:"backup_codes"` } // Setup2FA 初始化2FA设置 func Setup2FA(c *gin.Context) { userId := c.GetInt("id") // 检查用户是否已经启用2FA existing, err := model.GetTwoFAByUserId(userId) if err != nil { common.ApiError(c, err) return } if existing != nil && existing.IsEnabled { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "用户已启用2FA,请先禁用后重新设置", }) return } // 如果存在已禁用的2FA记录,先删除它 if existing != nil && !existing.IsEnabled { if err := existing.Delete(); err != nil { common.ApiError(c, err) return } existing = nil // 重置为nil,后续将创建新记录 } // 获取用户信息 user, err := model.GetUserById(userId, false) if err != nil { common.ApiError(c, err) return } // 生成TOTP密钥 key, err := common.GenerateTOTPSecret(user.Username) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "生成2FA密钥失败", }) common.SysLog("生成TOTP密钥失败: " + err.Error()) return } // 生成备用码 backupCodes, err := common.GenerateBackupCodes() if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "生成备用码失败", }) common.SysLog("生成备用码失败: " + err.Error()) return } // 生成二维码数据 qrCodeData := common.GenerateQRCodeData(key.Secret(), user.Username) // 创建或更新2FA记录(暂未启用) twoFA := &model.TwoFA{ UserId: userId, Secret: key.Secret(), IsEnabled: false, } if existing != nil { // 更新现有记录 twoFA.Id = existing.Id err = twoFA.Update() } else { // 创建新记录 err = twoFA.Create() } if err != nil { common.ApiError(c, err) return } // 创建备用码记录 if err := model.CreateBackupCodes(userId, backupCodes); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "保存备用码失败", }) common.SysLog("保存备用码失败: " + err.Error()) return } // 记录操作日志 model.RecordLog(userId, model.LogTypeSystem, "开始设置两步验证") c.JSON(http.StatusOK, gin.H{ "success": true, "message": "2FA设置初始化成功,请使用认证器扫描二维码并输入验证码完成设置", "data": Setup2FAResponse{ Secret: key.Secret(), QRCodeData: qrCodeData, BackupCodes: backupCodes, }, }) } // Enable2FA 启用2FA func Enable2FA(c *gin.Context) { var req Setup2FARequest if err := c.ShouldBindJSON(&req); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "参数错误", }) return } userId := c.GetInt("id") // 获取2FA记录 twoFA, err := model.GetTwoFAByUserId(userId) if err != nil { common.ApiError(c, err) return } if twoFA == nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "请先完成2FA初始化设置", }) return } if twoFA.IsEnabled { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "2FA已经启用", }) return } // 验证TOTP验证码 cleanCode, err := common.ValidateNumericCode(req.Code) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } if !common.ValidateTOTPCode(twoFA.Secret, cleanCode) { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "验证码或备用码错误,请重试", }) return } // 启用2FA if err := twoFA.Enable(); err != nil { common.ApiError(c, err) return } // 记录操作日志 model.RecordLog(userId, model.LogTypeSystem, "成功启用两步验证") c.JSON(http.StatusOK, gin.H{ "success": true, "message": "两步验证启用成功", }) } // Disable2FA 禁用2FA func Disable2FA(c *gin.Context) { var req Verify2FARequest if err := c.ShouldBindJSON(&req); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "参数错误", }) return } userId := c.GetInt("id") // 获取2FA记录 twoFA, err := model.GetTwoFAByUserId(userId) if err != nil { common.ApiError(c, err) return } if twoFA == nil || !twoFA.IsEnabled { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "用户未启用2FA", }) return } // 验证TOTP验证码或备用码 cleanCode, err := common.ValidateNumericCode(req.Code) isValidTOTP := false isValidBackup := false if err == nil { // 尝试验证TOTP isValidTOTP, _ = twoFA.ValidateTOTPAndUpdateUsage(cleanCode) } if !isValidTOTP { // 尝试验证备用码 isValidBackup, err = twoFA.ValidateBackupCodeAndUpdateUsage(req.Code) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } } if !isValidTOTP && !isValidBackup { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "验证码或备用码错误,请重试", }) return } // 禁用2FA if err := model.DisableTwoFA(userId); err != nil { common.ApiError(c, err) return } // 记录操作日志 model.RecordLog(userId, model.LogTypeSystem, "禁用两步验证") c.JSON(http.StatusOK, gin.H{ "success": true, "message": "两步验证已禁用", }) } // Get2FAStatus 获取用户2FA状态 func Get2FAStatus(c *gin.Context) { userId := c.GetInt("id") twoFA, err := model.GetTwoFAByUserId(userId) if err != nil { common.ApiError(c, err) return } status := map[string]interface{}{ "enabled": false, "locked": false, } if twoFA != nil { status["enabled"] = twoFA.IsEnabled status["locked"] = twoFA.IsLocked() if twoFA.IsEnabled { // 获取剩余备用码数量 backupCount, err := model.GetUnusedBackupCodeCount(userId) if err != nil { common.SysLog("获取备用码数量失败: " + err.Error()) } else { status["backup_codes_remaining"] = backupCount } } } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", "data": status, }) } // RegenerateBackupCodes 重新生成备用码 func RegenerateBackupCodes(c *gin.Context) { var req Verify2FARequest if err := c.ShouldBindJSON(&req); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "参数错误", }) return } userId := c.GetInt("id") // 获取2FA记录 twoFA, err := model.GetTwoFAByUserId(userId) if err != nil { common.ApiError(c, err) return } if twoFA == nil || !twoFA.IsEnabled { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "用户未启用2FA", }) return } // 验证TOTP验证码 cleanCode, err := common.ValidateNumericCode(req.Code) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } valid, err := twoFA.ValidateTOTPAndUpdateUsage(cleanCode) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } if !valid { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "验证码或备用码错误,请重试", }) return } // 生成新的备用码 backupCodes, err := common.GenerateBackupCodes() if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "生成备用码失败", }) common.SysLog("生成备用码失败: " + err.Error()) return } // 保存新的备用码 if err := model.CreateBackupCodes(userId, backupCodes); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "保存备用码失败", }) common.SysLog("保存备用码失败: " + err.Error()) return } // 记录操作日志 model.RecordLog(userId, model.LogTypeSystem, "重新生成两步验证备用码") c.JSON(http.StatusOK, gin.H{ "success": true, "message": "备用码重新生成成功", "data": map[string]interface{}{ "backup_codes": backupCodes, }, }) } // Verify2FALogin 登录时验证2FA func Verify2FALogin(c *gin.Context) { var req Verify2FARequest if err := c.ShouldBindJSON(&req); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "参数错误", }) return } // 从会话中获取pending用户信息 session := sessions.Default(c) pendingUserId := session.Get("pending_user_id") if pendingUserId == nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "会话已过期,请重新登录", }) return } userId, ok := pendingUserId.(int) if !ok { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "会话数据无效,请重新登录", }) return } // 获取用户信息 user, err := model.GetUserById(userId, false) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "用户不存在", }) return } // 获取2FA记录 twoFA, err := model.GetTwoFAByUserId(user.Id) if err != nil { common.ApiError(c, err) return } if twoFA == nil || !twoFA.IsEnabled { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "用户未启用2FA", }) return } // 验证TOTP验证码或备用码 cleanCode, err := common.ValidateNumericCode(req.Code) isValidTOTP := false isValidBackup := false if err == nil { // 尝试验证TOTP isValidTOTP, _ = twoFA.ValidateTOTPAndUpdateUsage(cleanCode) } if !isValidTOTP { // 尝试验证备用码 isValidBackup, err = twoFA.ValidateBackupCodeAndUpdateUsage(req.Code) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } } if !isValidTOTP && !isValidBackup { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "验证码或备用码错误,请重试", }) return } // 2FA验证成功,清理pending会话信息并完成登录 session.Delete("pending_username") session.Delete("pending_user_id") session.Save() setupLogin(user, c) } // Admin2FAStats 管理员获取2FA统计信息 func Admin2FAStats(c *gin.Context) { stats, err := model.GetTwoFAStats() if err != nil { common.ApiError(c, err) return } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", "data": stats, }) } // AdminDisable2FA 管理员强制禁用用户2FA func AdminDisable2FA(c *gin.Context) { userIdStr := c.Param("id") userId, err := strconv.Atoi(userIdStr) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "用户ID格式错误", }) return } // 检查目标用户权限 targetUser, err := model.GetUserById(userId, false) if err != nil { common.ApiError(c, err) return } myRole := c.GetInt("role") if myRole <= targetUser.Role && myRole != common.RoleRootUser { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无权操作同级或更高级用户的2FA设置", }) return } // 禁用2FA if err := model.DisableTwoFA(userId); err != nil { if errors.Is(err, model.ErrTwoFANotEnabled) { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "用户未启用2FA", }) return } common.ApiError(c, err) return } // 记录操作日志 adminId := c.GetInt("id") model.RecordLog(userId, model.LogTypeManage, fmt.Sprintf("管理员(ID:%d)强制禁用了用户的两步验证", adminId)) c.JSON(http.StatusOK, gin.H{ "success": true, "message": "用户2FA已被强制禁用", }) }