go_blog/models/db.go

package models

import (
	"crypto/rand"
	"log"
	"os"
	"path/filepath"

	"go_blog/config"

	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"github.com/glebarez/sqlite"
)

// DB is the global database connection, initialized by InitDB.
var DB *gorm.DB

// adminPasswordAlphabet avoids visually ambiguous characters (no l, I, O, 0,
// 1) and is used to generate the first-run admin password.
const adminPasswordAlphabet = "abcdefghjkmnpqrstuvwxyzABCDEFGHJKMNPQRSTUVWXYZ23456789" + "^$*+?%"

// randomAdminPassword returns a crypto-random 16-character first-run admin
// password (SECURITY_TODO #12: no more hardcoded admin/admin).
func randomAdminPassword() string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		log.Fatalf("Failed to generate admin password: %v", err)
	}
	out := make([]byte, len(b))
	for i, v := range b {
		out[i] = adminPasswordAlphabet[int(v)%len(adminPasswordAlphabet)]
	}
	return string(out)
}

// InitDB opens the database connection, runs migrations, and seeds the admin user.
func InitDB(cfg *config.Config) *gorm.DB {
	// Ensure the storage path exists.
	if err := os.MkdirAll(cfg.Path, 0755); err != nil {
		log.Fatalf("Failed to create storage directory %s: %v", cfg.Path, err)
	}

	var dialector gorm.Dialector

	switch cfg.Database.Type {
	case "mysql":
		if cfg.Database.DSN == "" {
			log.Fatalf("Database DSN is required when type is 'mysql'. Please set it in your config file.")
		}
		dialector = mysql.Open(cfg.Database.DSN)
	default:
		dbPath := filepath.Join(cfg.Path, "blog.db")
		dialector = sqlite.Open(dbPath)
	}

	db, err := gorm.Open(dialector, &gorm.Config{
		Logger: logger.Default.LogMode(logger.Warn),
	})
	if err != nil {
		log.Fatalf("Failed to connect to database: %v", err)
	}

	// Auto-migrate tables (idempotent).
	if err := db.AutoMigrate(&User{}, &Article{}, &SiteSetting{}, &UploadConfig{}, &UploadFileType{}, &DownloadBaseURL{}, &Attachment{}, &Comment{}, &CommentConfig{}, &ArticleView{}, &NavLink{}, &Tag{}, &ArticleTag{}); err != nil {
		log.Fatalf("Failed to auto-migrate database: %v", err)
	}

	// Seed site platform configuration on first run.
	seedSiteSettings(db)
	seedUploadConfig(db)
	seedUploadFileTypes(db)
	seedCommentConfig(db)

	// First-run seed: create admin user if no users exist.
	var count int64
	db.Model(&User{}).Count(&count)
	if count == 0 {
		admin := &User{
			Username:    "admin",
			DisplayName: "Administrator",
			Gender:      "other",
			Status:      StatusNormal,
			Role:        RoleAdmin,
		}
		adminPassword := randomAdminPassword()
		if err := admin.SetPassword(adminPassword); err != nil {
			log.Fatalf("Failed to hash admin password: %v", err)
		}
		if err := db.Create(admin).Error; err != nil {
			log.Fatalf("Failed to create admin user: %v", err)
		}
		// SECURITY_TODO #12: the first-run password is crypto-random and
		// printed exactly once — copy it now; it cannot be recovered later.
		log.Println("==============================================")
		log.Println("  First run: created default admin user.")
		log.Println("  Username: admin")
		log.Println("  Password: " + adminPassword)
		log.Println("  This password is shown ONCE. Change it after login!")
		log.Println("==============================================")
	}

	// Migration fix: always set admin role on the admin user.
	result := db.Model(&User{}).Where("username = ?", "admin").Update("role", RoleAdmin)
	if result.RowsAffected > 0 {
		log.Printf("Migration: set admin role for existing admin user (rows affected: %d)", result.RowsAffected)
	}

	DB = db
	return db
}