package model import ( "errors" "fmt" "time" "github.com/QuantumNous/new-api/common" "gorm.io/gorm" ) var ErrTwoFANotEnabled = errors.New("用户未启用2FA") // 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"` 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"` 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不能为空") } // 使用事务确保原子性 return DB.Transaction(func(tx *gorm.DB) error { // 同时删除相关的备用码记录(硬删除) if err := tx.Unscoped().Where("user_id = ?", t.UserId).Delete(&TwoFABackupCode{}).Error; err != nil { return err } // 硬删除2FA记录 return tx.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 { return DB.Transaction(func(tx *gorm.DB) error { // 先删除现有的备用码 if err := tx.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 := tx.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 ErrTwoFANotEnabled } // 删除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.SysLog("更新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.SysLog("更新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.SysLog("更新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.SysLog("更新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 }