88 lines
3.3 KiB
Go
88 lines
3.3 KiB
Go
// Filename: internal/db/seeder/seeder.go
|
|
package seeder
|
|
|
|
import (
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"gemini-balancer/internal/crypto"
|
|
"gemini-balancer/internal/models"
|
|
"gemini-balancer/internal/repository"
|
|
|
|
"github.com/sirupsen/logrus"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
// RunSeeder now requires the crypto service to create the initial admin token securely.
|
|
func RunSeeder(db *gorm.DB, cryptoService *crypto.Service, logger *logrus.Logger) {
|
|
log := logger.WithField("component", "seeder")
|
|
log.Info("Running database seeder...")
|
|
// [REFACTORED] Admin token seeding is now crypto-aware.
|
|
var count int64
|
|
db.Model(&models.AuthToken{}).Where("is_admin = ?", true).Count(&count)
|
|
if count == 0 {
|
|
log.Info("No admin token found, attempting to seed one...")
|
|
const adminTokenPlaintext = "admin-secret-token" // The default token
|
|
// 1. Encrypt and Hash the token
|
|
encryptedToken, err := cryptoService.Encrypt(adminTokenPlaintext)
|
|
if err != nil {
|
|
log.Fatalf("FATAL: Failed to encrypt default admin token during seeding: %v. Server cannot start.", err)
|
|
return
|
|
}
|
|
hash := sha256.Sum256([]byte(adminTokenPlaintext))
|
|
tokenHash := hex.EncodeToString(hash[:])
|
|
// 2. Use the repository to seed the token
|
|
// Note: We create a temporary repository instance here just for the seeder.
|
|
repo := repository.NewAuthTokenRepository(db, cryptoService, logger)
|
|
if err := repo.SeedAdminToken(encryptedToken, tokenHash); err != nil {
|
|
log.Warnf("Failed to seed admin token using repository: %v", err)
|
|
} else {
|
|
log.Infof("Default admin token has been seeded successfully. Please use '%s' for your first login.", adminTokenPlaintext)
|
|
}
|
|
} else {
|
|
log.Info("Admin token already exists, seeder skipped.")
|
|
}
|
|
|
|
// This functionality should be replaced by a proper user/token management UI in the future.
|
|
linkAllKeysToDefaultGroup(db, log)
|
|
}
|
|
|
|
// linkAllKeysToDefaultGroup ensures every key belongs to at least one group.
|
|
func linkAllKeysToDefaultGroup(db *gorm.DB, logger *logrus.Entry) {
|
|
logger.Info("Linking existing API keys to the default group as a fallback...")
|
|
// 1. Find a default group (the first one for simplicity)
|
|
var defaultGroup models.KeyGroup
|
|
if err := db.Order("id asc").First(&defaultGroup).Error; err != nil {
|
|
logger.Warnf("Seeder: Could not find a default key group to link keys to: %v", err)
|
|
return
|
|
}
|
|
// 2. Find all "orphan keys" that don't belong to any group
|
|
var orphanKeys []*models.APIKey
|
|
err := db.Raw(`
|
|
SELECT * FROM api_keys
|
|
WHERE id NOT IN (SELECT DISTINCT api_key_id FROM group_api_key_mappings)
|
|
AND deleted_at IS NULL
|
|
`).Scan(&orphanKeys).Error
|
|
if err != nil {
|
|
logger.Errorf("Seeder: Failed to query for orphan keys: %v", err)
|
|
return
|
|
}
|
|
if len(orphanKeys) == 0 {
|
|
logger.Info("Seeder: No orphan API keys found to link.")
|
|
return
|
|
}
|
|
// 3. Create GroupAPIKeyMapping records manually
|
|
logger.Infof("Seeder: Found %d orphan keys. Creating mappings for them in group '%s' (ID: %d)...", len(orphanKeys), defaultGroup.Name, defaultGroup.ID)
|
|
var newMappings []models.GroupAPIKeyMapping
|
|
for _, key := range orphanKeys {
|
|
newMappings = append(newMappings, models.GroupAPIKeyMapping{
|
|
KeyGroupID: defaultGroup.ID,
|
|
APIKeyID: key.ID,
|
|
})
|
|
}
|
|
if err := db.Create(&newMappings).Error; err != nil {
|
|
logger.Errorf("Seeder: Failed to create key mappings for orphan keys: %v", err)
|
|
} else {
|
|
logger.Info("Successfully created mappings for orphan API keys.")
|
|
}
|
|
}
|