6748b006b7
This commit refactors the logging mechanism across the application by replacing direct logger calls with a centralized logging approach using the `common` package. Key changes include: - Replaced instances of `logger.SysLog` and `logger.FatalLog` with `common.SysLog` and `common.FatalLog` for consistent logging practices. - Updated resource initialization error handling to utilize the new logging structure, enhancing maintainability and readability. - Minor adjustments to improve code clarity and organization throughout various modules. This change aims to streamline logging and improve the overall architecture of the codebase.
323 lines
8.0 KiB
Go
323 lines
8.0 KiB
Go
package model
|
||
|
||
import (
|
||
"errors"
|
||
"fmt"
|
||
"one-api/common"
|
||
"time"
|
||
|
||
"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" 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不能为空")
|
||
}
|
||
|
||
// 使用事务确保原子性
|
||
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
|
||
}
|