feat: implement two-factor authentication (2FA) support with user login and settings integration
This commit is contained in:
Seefs
547
controller/twofa.go
Normal file
547
controller/twofa.go
Normal file
@@ -0,0 +1,547 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"one-api/common"
|
||||
"one-api/model"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"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.SysError("生成TOTP密钥失败: " + err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// 生成备用码
|
||||
backupCodes, err := common.GenerateBackupCodes()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "生成备用码失败",
|
||||
})
|
||||
common.SysError("生成备用码失败: " + 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.SysError("保存备用码失败: " + 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.SysError("获取备用码数量失败: " + 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.SysError("生成备用码失败: " + err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// 保存新的备用码
|
||||
if err := model.CreateBackupCodes(userId, backupCodes); err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "保存备用码失败",
|
||||
})
|
||||
common.SysError("保存备用码失败: " + 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 := pendingUserId.(int)
|
||||
|
||||
// 获取用户信息
|
||||
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 strings.Contains(err.Error(), "未启用2FA") {
|
||||
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已被强制禁用",
|
||||
})
|
||||
}
|
||||
@@ -62,6 +62,32 @@ func Login(c *gin.Context) {
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// 检查是否启用2FA
|
||||
if model.IsTwoFAEnabled(user.Id) {
|
||||
// 设置pending session，等待2FA验证
|
||||
session := sessions.Default(c)
|
||||
session.Set("pending_username", user.Username)
|
||||
session.Set("pending_user_id", user.Id)
|
||||
err := session.Save()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "无法保存会话信息，请重试",
|
||||
"success": false,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "请输入两步验证码",
|
||||
"success": true,
|
||||
"data": map[string]interface{}{
|
||||
"require_2fa": true,
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
setupLogin(&user, c)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user