101 lines
2.1 KiB
Go
101 lines
2.1 KiB
Go
package database
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"anticaptcha/internal/config"
|
|
"anticaptcha/internal/model"
|
|
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/sqlite"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
)
|
|
|
|
var DB *gorm.DB
|
|
|
|
func Init(cfg *config.DatabaseConfig) error {
|
|
var err error
|
|
var gormConfig *gorm.Config
|
|
|
|
// 生产环境禁用日志
|
|
gormConfig = &gorm.Config{
|
|
Logger: logger.Default.LogMode(logger.Silent),
|
|
}
|
|
|
|
switch cfg.Type {
|
|
case "sqlite":
|
|
DB, err = gorm.Open(sqlite.Open(cfg.SQLite.Path), gormConfig)
|
|
case "mysql":
|
|
dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4&parseTime=True&loc=Local",
|
|
cfg.User, cfg.Password, cfg.Host, cfg.Port, cfg.Database)
|
|
DB, err = gorm.Open(mysql.Open(dsn), gormConfig)
|
|
default:
|
|
return fmt.Errorf("不支持的数据库类型: %s", cfg.Type)
|
|
}
|
|
|
|
if err != nil {
|
|
return fmt.Errorf("数据库连接失败: %v", err)
|
|
}
|
|
|
|
// 自动迁移
|
|
if err = DB.AutoMigrate(
|
|
&model.User{},
|
|
&model.RegistrationCode{},
|
|
&model.EndpointCost{},
|
|
&model.Config{},
|
|
); err != nil {
|
|
return fmt.Errorf("数据库迁移失败: %v", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func Close() error {
|
|
if DB != nil {
|
|
sqlDB, err := DB.DB()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return sqlDB.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// 确保 admin 用户存在（仅在安装时调用）
|
|
func EnsureAdminUser(username, password string) error {
|
|
var count int64
|
|
DB.Model(&model.User{}).Where("role = ?", model.RoleAdmin).Count(&count)
|
|
if count > 0 {
|
|
return nil
|
|
}
|
|
|
|
admin := &model.User{
|
|
Username: username,
|
|
HashedPassword: HashPassword(password),
|
|
Role: model.RoleAdmin,
|
|
Balance: 1000000,
|
|
}
|
|
|
|
return DB.Create(admin).Error
|
|
}
|
|
|
|
// 验证用户登录
|
|
func ValidateUser(username, password string) (*model.User, error) {
|
|
var user model.User
|
|
if err := DB.Where("username = ?", username).First(&user).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if user.HashedPassword != password {
|
|
return nil, fmt.Errorf("密码错误")
|
|
}
|
|
|
|
return &user, nil
|
|
}
|
|
|
|
func HashPassword(password string) string {
|
|
// 简化处理，实际应使用 bcrypt
|
|
// 后续补充完整实现
|
|
return password
|
|
} |