316 lines
7.8 KiB
Go
316 lines
7.8 KiB
Go
package model
|
||
|
||
import (
|
||
"errors"
|
||
"fmt"
|
||
"one-api/common"
|
||
"time"
|
||
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
// TwoFA 用户2FA设置表
|
||
type TwoFA struct {
|
||
Id int `json:"id" gorm:"primaryKey"`
|
||
UserId int `json:"user_id" gorm:"unique;not null;index"`
|
||
Secret string `json:"-" gorm:"type:varchar(255);not null"` // TOTP密钥，不返回给前端
|
||
IsEnabled bool `json:"is_enabled" gorm:"default:false"`
|
||
FailedAttempts int `json:"failed_attempts" gorm:"default:0"`
|
||
LockedUntil *time.Time `json:"locked_until,omitempty"`
|
||
LastUsedAt *time.Time `json:"last_used_at,omitempty"`
|
||
CreatedAt time.Time `json:"created_at"`
|
||
UpdatedAt time.Time `json:"updated_at"`
|
||
DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
|
||
}
|
||
|
||
// TwoFABackupCode 备用码使用记录表
|
||
type TwoFABackupCode struct {
|
||
Id int `json:"id" gorm:"primaryKey"`
|
||
UserId int `json:"user_id" gorm:"not null;index"`
|
||
CodeHash string `json:"-" gorm:"type:varchar(255);not null"` // 备用码哈希
|
||
IsUsed bool `json:"is_used" gorm:"default:false"`
|
||
UsedAt *time.Time `json:"used_at,omitempty"`
|
||
CreatedAt time.Time `json:"created_at"`
|
||
DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
|
||
}
|
||
|
||
// GetTwoFAByUserId 根据用户ID获取2FA设置
|
||
func GetTwoFAByUserId(userId int) (*TwoFA, error) {
|
||
if userId == 0 {
|
||
return nil, errors.New("用户ID不能为空")
|
||
}
|
||
|
||
var twoFA TwoFA
|
||
err := DB.Where("user_id = ?", userId).First(&twoFA).Error
|
||
if err != nil {
|
||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||
return nil, nil // 返回nil表示未设置2FA
|
||
}
|
||
return nil, err
|
||
}
|
||
|
||
return &twoFA, nil
|
||
}
|
||
|
||
// IsTwoFAEnabled 检查用户是否启用了2FA
|
||
func IsTwoFAEnabled(userId int) bool {
|
||
twoFA, err := GetTwoFAByUserId(userId)
|
||
if err != nil || twoFA == nil {
|
||
return false
|
||
}
|
||
return twoFA.IsEnabled
|
||
}
|
||
|
||
// CreateTwoFA 创建2FA设置
|
||
func (t *TwoFA) Create() error {
|
||
// 检查用户是否已存在2FA设置
|
||
existing, err := GetTwoFAByUserId(t.UserId)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if existing != nil {
|
||
return errors.New("用户已存在2FA设置")
|
||
}
|
||
|
||
// 验证用户存在
|
||
var user User
|
||
if err := DB.First(&user, t.UserId).Error; err != nil {
|
||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||
return errors.New("用户不存在")
|
||
}
|
||
return err
|
||
}
|
||
|
||
return DB.Create(t).Error
|
||
}
|
||
|
||
// Update 更新2FA设置
|
||
func (t *TwoFA) Update() error {
|
||
if t.Id == 0 {
|
||
return errors.New("2FA记录ID不能为空")
|
||
}
|
||
return DB.Save(t).Error
|
||
}
|
||
|
||
// Delete 删除2FA设置
|
||
func (t *TwoFA) Delete() error {
|
||
if t.Id == 0 {
|
||
return errors.New("2FA记录ID不能为空")
|
||
}
|
||
|
||
// 同时删除相关的备用码记录（硬删除）
|
||
if err := DB.Unscoped().Where("user_id = ?", t.UserId).Delete(&TwoFABackupCode{}).Error; err != nil {
|
||
return err
|
||
}
|
||
|
||
// 硬删除2FA记录
|
||
return DB.Unscoped().Delete(t).Error
|
||
}
|
||
|
||
// ResetFailedAttempts 重置失败尝试次数
|
||
func (t *TwoFA) ResetFailedAttempts() error {
|
||
t.FailedAttempts = 0
|
||
t.LockedUntil = nil
|
||
return t.Update()
|
||
}
|
||
|
||
// IncrementFailedAttempts 增加失败尝试次数
|
||
func (t *TwoFA) IncrementFailedAttempts() error {
|
||
t.FailedAttempts++
|
||
|
||
// 检查是否需要锁定
|
||
if t.FailedAttempts >= common.MaxFailAttempts {
|
||
lockUntil := time.Now().Add(time.Duration(common.LockoutDuration) * time.Second)
|
||
t.LockedUntil = &lockUntil
|
||
}
|
||
|
||
return t.Update()
|
||
}
|
||
|
||
// IsLocked 检查账户是否被锁定
|
||
func (t *TwoFA) IsLocked() bool {
|
||
if t.LockedUntil == nil {
|
||
return false
|
||
}
|
||
return time.Now().Before(*t.LockedUntil)
|
||
}
|
||
|
||
// CreateBackupCodes 创建备用码
|
||
func CreateBackupCodes(userId int, codes []string) error {
|
||
// 先删除现有的备用码
|
||
if err := DB.Where("user_id = ?", userId).Delete(&TwoFABackupCode{}).Error; err != nil {
|
||
return err
|
||
}
|
||
|
||
// 创建新的备用码记录
|
||
for _, code := range codes {
|
||
hashedCode, err := common.HashBackupCode(code)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
backupCode := TwoFABackupCode{
|
||
UserId: userId,
|
||
CodeHash: hashedCode,
|
||
IsUsed: false,
|
||
}
|
||
|
||
if err := DB.Create(&backupCode).Error; err != nil {
|
||
return err
|
||
}
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// ValidateBackupCode 验证并使用备用码
|
||
func ValidateBackupCode(userId int, code string) (bool, error) {
|
||
if !common.ValidateBackupCode(code) {
|
||
return false, errors.New("验证码或备用码不正确")
|
||
}
|
||
|
||
normalizedCode := common.NormalizeBackupCode(code)
|
||
|
||
// 查找未使用的备用码
|
||
var backupCodes []TwoFABackupCode
|
||
if err := DB.Where("user_id = ? AND is_used = false", userId).Find(&backupCodes).Error; err != nil {
|
||
return false, err
|
||
}
|
||
|
||
// 验证备用码
|
||
for _, bc := range backupCodes {
|
||
if common.ValidatePasswordAndHash(normalizedCode, bc.CodeHash) {
|
||
// 标记为已使用
|
||
now := time.Now()
|
||
bc.IsUsed = true
|
||
bc.UsedAt = &now
|
||
|
||
if err := DB.Save(&bc).Error; err != nil {
|
||
return false, err
|
||
}
|
||
|
||
return true, nil
|
||
}
|
||
}
|
||
|
||
return false, nil
|
||
}
|
||
|
||
// GetUnusedBackupCodeCount 获取未使用的备用码数量
|
||
func GetUnusedBackupCodeCount(userId int) (int, error) {
|
||
var count int64
|
||
err := DB.Model(&TwoFABackupCode{}).Where("user_id = ? AND is_used = false", userId).Count(&count).Error
|
||
return int(count), err
|
||
}
|
||
|
||
// DisableTwoFA 禁用用户的2FA
|
||
func DisableTwoFA(userId int) error {
|
||
twoFA, err := GetTwoFAByUserId(userId)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if twoFA == nil {
|
||
return errors.New("用户未启用2FA")
|
||
}
|
||
|
||
// 删除2FA设置和备用码
|
||
return twoFA.Delete()
|
||
}
|
||
|
||
// EnableTwoFA 启用2FA
|
||
func (t *TwoFA) Enable() error {
|
||
t.IsEnabled = true
|
||
t.FailedAttempts = 0
|
||
t.LockedUntil = nil
|
||
return t.Update()
|
||
}
|
||
|
||
// ValidateTOTPAndUpdateUsage 验证TOTP并更新使用记录
|
||
func (t *TwoFA) ValidateTOTPAndUpdateUsage(code string) (bool, error) {
|
||
// 检查是否被锁定
|
||
if t.IsLocked() {
|
||
return false, fmt.Errorf("账户已被锁定，请在%v后重试", t.LockedUntil.Format("2006-01-02 15:04:05"))
|
||
}
|
||
|
||
// 验证TOTP码
|
||
if !common.ValidateTOTPCode(t.Secret, code) {
|
||
// 增加失败次数
|
||
if err := t.IncrementFailedAttempts(); err != nil {
|
||
common.SysError("更新2FA失败次数失败: " + err.Error())
|
||
}
|
||
return false, nil
|
||
}
|
||
|
||
// 验证成功，重置失败次数并更新最后使用时间
|
||
now := time.Now()
|
||
t.FailedAttempts = 0
|
||
t.LockedUntil = nil
|
||
t.LastUsedAt = &now
|
||
|
||
if err := t.Update(); err != nil {
|
||
common.SysError("更新2FA使用记录失败: " + err.Error())
|
||
}
|
||
|
||
return true, nil
|
||
}
|
||
|
||
// ValidateBackupCodeAndUpdateUsage 验证备用码并更新使用记录
|
||
func (t *TwoFA) ValidateBackupCodeAndUpdateUsage(code string) (bool, error) {
|
||
// 检查是否被锁定
|
||
if t.IsLocked() {
|
||
return false, fmt.Errorf("账户已被锁定，请在%v后重试", t.LockedUntil.Format("2006-01-02 15:04:05"))
|
||
}
|
||
|
||
// 验证备用码
|
||
valid, err := ValidateBackupCode(t.UserId, code)
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
|
||
if !valid {
|
||
// 增加失败次数
|
||
if err := t.IncrementFailedAttempts(); err != nil {
|
||
common.SysError("更新2FA失败次数失败: " + err.Error())
|
||
}
|
||
return false, nil
|
||
}
|
||
|
||
// 验证成功，重置失败次数并更新最后使用时间
|
||
now := time.Now()
|
||
t.FailedAttempts = 0
|
||
t.LockedUntil = nil
|
||
t.LastUsedAt = &now
|
||
|
||
if err := t.Update(); err != nil {
|
||
common.SysError("更新2FA使用记录失败: " + err.Error())
|
||
}
|
||
|
||
return true, nil
|
||
}
|
||
|
||
// GetTwoFAStats 获取2FA统计信息（管理员使用）
|
||
func GetTwoFAStats() (map[string]interface{}, error) {
|
||
var totalUsers, enabledUsers int64
|
||
|
||
// 总用户数
|
||
if err := DB.Model(&User{}).Count(&totalUsers).Error; err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 启用2FA的用户数
|
||
if err := DB.Model(&TwoFA{}).Where("is_enabled = true").Count(&enabledUsers).Error; err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
enabledRate := float64(0)
|
||
if totalUsers > 0 {
|
||
enabledRate = float64(enabledUsers) / float64(totalUsers) * 100
|
||
}
|
||
|
||
return map[string]interface{}{
|
||
"total_users": totalUsers,
|
||
"enabled_users": enabledUsers,
|
||
"enabled_rate": fmt.Sprintf("%.1f%%", enabledRate),
|
||
}, nil
|
||
}
|