package setup import ( "context" "crypto/rand" "encoding/hex" "fmt" "log" "os" "strconv" "time" "github.com/redis/go-redis/v9" "golang.org/x/crypto/bcrypt" "gorm.io/driver/postgres" "gorm.io/gorm" "gopkg.in/yaml.v3" ) // Config paths const ( ConfigFile = "config.yaml" EnvFile = ".env" ) // SetupConfig holds the setup configuration type SetupConfig struct { Database DatabaseConfig `json:"database" yaml:"database"` Redis RedisConfig `json:"redis" yaml:"redis"` Admin AdminConfig `json:"admin" yaml:"-"` // Not stored in config file Server ServerConfig `json:"server" yaml:"server"` JWT JWTConfig `json:"jwt" yaml:"jwt"` Timezone string `json:"timezone" yaml:"timezone"` // e.g. "Asia/Shanghai", "UTC" } type DatabaseConfig struct { Host string `json:"host" yaml:"host"` Port int `json:"port" yaml:"port"` User string `json:"user" yaml:"user"` Password string `json:"password" yaml:"password"` DBName string `json:"dbname" yaml:"dbname"` SSLMode string `json:"sslmode" yaml:"sslmode"` } type RedisConfig struct { Host string `json:"host" yaml:"host"` Port int `json:"port" yaml:"port"` Password string `json:"password" yaml:"password"` DB int `json:"db" yaml:"db"` } type AdminConfig struct { Email string `json:"email"` Password string `json:"password"` } type ServerConfig struct { Host string `json:"host" yaml:"host"` Port int `json:"port" yaml:"port"` Mode string `json:"mode" yaml:"mode"` } type JWTConfig struct { Secret string `json:"secret" yaml:"secret"` ExpireHour int `json:"expire_hour" yaml:"expire_hour"` } // NeedsSetup checks if the system needs initial setup // Uses multiple checks to prevent attackers from forcing re-setup by deleting config func NeedsSetup() bool { // Check 1: Config file must not exist if _, err := os.Stat(ConfigFile); !os.IsNotExist(err) { return false // Config exists, no setup needed } // Check 2: Installation lock file (harder to bypass) lockFile := ".installed" if _, err := os.Stat(lockFile); !os.IsNotExist(err) { return false // Lock file exists, already installed } return true } // TestDatabaseConnection tests the database connection func TestDatabaseConnection(cfg *DatabaseConfig) error { dsn := fmt.Sprintf( "host=%s port=%d user=%s password=%s dbname=%s sslmode=%s", cfg.Host, cfg.Port, cfg.User, cfg.Password, cfg.DBName, cfg.SSLMode, ) db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{}) if err != nil { return fmt.Errorf("failed to connect: %w", err) } sqlDB, err := db.DB() if err != nil { return fmt.Errorf("failed to get db instance: %w", err) } defer sqlDB.Close() ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if err := sqlDB.PingContext(ctx); err != nil { return fmt.Errorf("ping failed: %w", err) } return nil } // TestRedisConnection tests the Redis connection func TestRedisConnection(cfg *RedisConfig) error { rdb := redis.NewClient(&redis.Options{ Addr: fmt.Sprintf("%s:%d", cfg.Host, cfg.Port), Password: cfg.Password, DB: cfg.DB, }) defer rdb.Close() ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if err := rdb.Ping(ctx).Err(); err != nil { return fmt.Errorf("ping failed: %w", err) } return nil } // Install performs the installation with the given configuration func Install(cfg *SetupConfig) error { // Security check: prevent re-installation if already installed if !NeedsSetup() { return fmt.Errorf("system is already installed, re-installation is not allowed") } // Generate JWT secret if not provided if cfg.JWT.Secret == "" { cfg.JWT.Secret = generateSecret(32) } // Test connections if err := TestDatabaseConnection(&cfg.Database); err != nil { return fmt.Errorf("database connection failed: %w", err) } if err := TestRedisConnection(&cfg.Redis); err != nil { return fmt.Errorf("redis connection failed: %w", err) } // Initialize database if err := initializeDatabase(cfg); err != nil { return fmt.Errorf("database initialization failed: %w", err) } // Create admin user if err := createAdminUser(cfg); err != nil { return fmt.Errorf("admin user creation failed: %w", err) } // Write config file if err := writeConfigFile(cfg); err != nil { return fmt.Errorf("config file creation failed: %w", err) } // Create installation lock file to prevent re-setup attacks if err := createInstallLock(); err != nil { return fmt.Errorf("failed to create install lock: %w", err) } return nil } // createInstallLock creates a lock file to prevent re-installation attacks func createInstallLock() error { lockFile := ".installed" content := fmt.Sprintf("installed_at=%s\n", time.Now().UTC().Format(time.RFC3339)) return os.WriteFile(lockFile, []byte(content), 0400) // Read-only for owner } func initializeDatabase(cfg *SetupConfig) error { dsn := fmt.Sprintf( "host=%s port=%d user=%s password=%s dbname=%s sslmode=%s", cfg.Database.Host, cfg.Database.Port, cfg.Database.User, cfg.Database.Password, cfg.Database.DBName, cfg.Database.SSLMode, ) db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{}) if err != nil { return err } sqlDB, err := db.DB() if err != nil { return err } defer sqlDB.Close() // Run auto-migration for all models return db.AutoMigrate( &User{}, &Group{}, &APIKey{}, &Account{}, &Proxy{}, &RedeemCode{}, &UsageLog{}, &UserSubscription{}, &Setting{}, ) } func createAdminUser(cfg *SetupConfig) error { dsn := fmt.Sprintf( "host=%s port=%d user=%s password=%s dbname=%s sslmode=%s", cfg.Database.Host, cfg.Database.Port, cfg.Database.User, cfg.Database.Password, cfg.Database.DBName, cfg.Database.SSLMode, ) db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{}) if err != nil { return err } sqlDB, err := db.DB() if err != nil { return err } defer sqlDB.Close() // Check if admin already exists var count int64 db.Model(&User{}).Where("role = ?", "admin").Count(&count) if count > 0 { return nil // Admin already exists } // Hash password hashedPassword, err := bcrypt.GenerateFromPassword([]byte(cfg.Admin.Password), bcrypt.DefaultCost) if err != nil { return err } // Create admin user admin := &User{ Email: cfg.Admin.Email, PasswordHash: string(hashedPassword), Role: "admin", Status: "active", Balance: 0, CreatedAt: time.Now(), UpdatedAt: time.Now(), } return db.Create(admin).Error } func writeConfigFile(cfg *SetupConfig) error { // Ensure timezone has a default value tz := cfg.Timezone if tz == "" { tz = "Asia/Shanghai" } // Prepare config for YAML (exclude sensitive data and admin config) yamlConfig := struct { Server ServerConfig `yaml:"server"` Database DatabaseConfig `yaml:"database"` Redis RedisConfig `yaml:"redis"` JWT struct { Secret string `yaml:"secret"` ExpireHour int `yaml:"expire_hour"` } `yaml:"jwt"` Default struct { GroupID uint `yaml:"group_id"` } `yaml:"default"` RateLimit struct { RequestsPerMinute int `yaml:"requests_per_minute"` BurstSize int `yaml:"burst_size"` } `yaml:"rate_limit"` Timezone string `yaml:"timezone"` }{ Server: cfg.Server, Database: cfg.Database, Redis: cfg.Redis, JWT: struct { Secret string `yaml:"secret"` ExpireHour int `yaml:"expire_hour"` }{ Secret: cfg.JWT.Secret, ExpireHour: cfg.JWT.ExpireHour, }, Default: struct { GroupID uint `yaml:"group_id"` }{ GroupID: 1, }, RateLimit: struct { RequestsPerMinute int `yaml:"requests_per_minute"` BurstSize int `yaml:"burst_size"` }{ RequestsPerMinute: 60, BurstSize: 10, }, Timezone: tz, } data, err := yaml.Marshal(&yamlConfig) if err != nil { return err } return os.WriteFile(ConfigFile, data, 0600) } func generateSecret(length int) string { bytes := make([]byte, length) rand.Read(bytes) return hex.EncodeToString(bytes) } // Minimal model definitions for migration (to avoid circular import) type User struct { ID uint `gorm:"primaryKey"` Email string `gorm:"uniqueIndex;not null"` PasswordHash string `gorm:"not null"` Role string `gorm:"default:user"` Status string `gorm:"default:active"` Balance float64 `gorm:"default:0"` CreatedAt time.Time UpdatedAt time.Time } type Group struct { ID uint `gorm:"primaryKey"` Name string `gorm:"uniqueIndex;not null"` Description string `gorm:"type:text"` RateMultiplier float64 `gorm:"default:1.0"` IsExclusive bool `gorm:"default:false"` Priority int `gorm:"default:0"` Status string `gorm:"default:active"` CreatedAt time.Time UpdatedAt time.Time } type APIKey struct { ID uint `gorm:"primaryKey"` UserID uint `gorm:"index;not null"` Key string `gorm:"uniqueIndex;not null"` Name string GroupID *uint Status string `gorm:"default:active"` CreatedAt time.Time UpdatedAt time.Time } type Account struct { ID uint `gorm:"primaryKey"` Platform string `gorm:"not null"` Type string `gorm:"not null"` Credentials string `gorm:"type:text"` Status string `gorm:"default:active"` Priority int `gorm:"default:0"` ProxyID *uint CreatedAt time.Time UpdatedAt time.Time } type Proxy struct { ID uint `gorm:"primaryKey"` Name string `gorm:"not null"` Protocol string `gorm:"not null"` Host string `gorm:"not null"` Port int `gorm:"not null"` Username string Password string Status string `gorm:"default:active"` CreatedAt time.Time UpdatedAt time.Time } type RedeemCode struct { ID uint `gorm:"primaryKey"` Code string `gorm:"uniqueIndex;not null"` Value float64 `gorm:"not null"` Status string `gorm:"default:unused"` UsedBy *uint UsedAt *time.Time ExpiresAt *time.Time CreatedAt time.Time } type UsageLog struct { ID uint `gorm:"primaryKey"` UserID uint `gorm:"index"` APIKeyID uint `gorm:"index"` AccountID *uint `gorm:"index"` Model string `gorm:"index"` InputTokens int OutputTokens int Cost float64 CreatedAt time.Time } type UserSubscription struct { ID uint `gorm:"primaryKey"` UserID uint `gorm:"index;not null"` GroupID uint `gorm:"index;not null"` Quota int64 Used int64 `gorm:"default:0"` Status string ExpiresAt *time.Time CreatedAt time.Time UpdatedAt time.Time } type Setting struct { ID uint `gorm:"primaryKey"` Key string `gorm:"uniqueIndex;not null"` Value string `gorm:"type:text"` CreatedAt time.Time UpdatedAt time.Time } func (User) TableName() string { return "users" } func (Group) TableName() string { return "groups" } func (APIKey) TableName() string { return "api_keys" } func (Account) TableName() string { return "accounts" } func (Proxy) TableName() string { return "proxies" } func (RedeemCode) TableName() string { return "redeem_codes" } func (UsageLog) TableName() string { return "usage_logs" } func (UserSubscription) TableName() string { return "user_subscriptions" } func (Setting) TableName() string { return "settings" } // ============================================================================= // Auto Setup for Docker Deployment // ============================================================================= // AutoSetupEnabled checks if auto setup is enabled via environment variable func AutoSetupEnabled() bool { val := os.Getenv("AUTO_SETUP") return val == "true" || val == "1" || val == "yes" } // getEnvOrDefault gets environment variable or returns default value func getEnvOrDefault(key, defaultValue string) string { if val := os.Getenv(key); val != "" { return val } return defaultValue } // getEnvIntOrDefault gets environment variable as int or returns default value func getEnvIntOrDefault(key string, defaultValue int) int { if val := os.Getenv(key); val != "" { if i, err := strconv.Atoi(val); err == nil { return i } } return defaultValue } // AutoSetupFromEnv performs automatic setup using environment variables // This is designed for Docker deployment where all config is passed via env vars func AutoSetupFromEnv() error { log.Println("Auto setup enabled, configuring from environment variables...") // Get timezone from TZ or TIMEZONE env var (TZ is standard for Docker) tz := getEnvOrDefault("TZ", "") if tz == "" { tz = getEnvOrDefault("TIMEZONE", "Asia/Shanghai") } // Build config from environment variables cfg := &SetupConfig{ Database: DatabaseConfig{ Host: getEnvOrDefault("DATABASE_HOST", "localhost"), Port: getEnvIntOrDefault("DATABASE_PORT", 5432), User: getEnvOrDefault("DATABASE_USER", "postgres"), Password: getEnvOrDefault("DATABASE_PASSWORD", ""), DBName: getEnvOrDefault("DATABASE_DBNAME", "sub2api"), SSLMode: getEnvOrDefault("DATABASE_SSLMODE", "disable"), }, Redis: RedisConfig{ Host: getEnvOrDefault("REDIS_HOST", "localhost"), Port: getEnvIntOrDefault("REDIS_PORT", 6379), Password: getEnvOrDefault("REDIS_PASSWORD", ""), DB: getEnvIntOrDefault("REDIS_DB", 0), }, Admin: AdminConfig{ Email: getEnvOrDefault("ADMIN_EMAIL", "admin@sub2api.local"), Password: getEnvOrDefault("ADMIN_PASSWORD", ""), }, Server: ServerConfig{ Host: getEnvOrDefault("SERVER_HOST", "0.0.0.0"), Port: getEnvIntOrDefault("SERVER_PORT", 8080), Mode: getEnvOrDefault("SERVER_MODE", "release"), }, JWT: JWTConfig{ Secret: getEnvOrDefault("JWT_SECRET", ""), ExpireHour: getEnvIntOrDefault("JWT_EXPIRE_HOUR", 24), }, Timezone: tz, } // Generate JWT secret if not provided if cfg.JWT.Secret == "" { cfg.JWT.Secret = generateSecret(32) log.Println("Generated JWT secret automatically") } // Generate admin password if not provided if cfg.Admin.Password == "" { cfg.Admin.Password = generateSecret(16) log.Printf("Generated admin password: %s", cfg.Admin.Password) log.Println("IMPORTANT: Save this password! It will not be shown again.") } // Test database connection log.Println("Testing database connection...") if err := TestDatabaseConnection(&cfg.Database); err != nil { return fmt.Errorf("database connection failed: %w", err) } log.Println("Database connection successful") // Test Redis connection log.Println("Testing Redis connection...") if err := TestRedisConnection(&cfg.Redis); err != nil { return fmt.Errorf("redis connection failed: %w", err) } log.Println("Redis connection successful") // Initialize database log.Println("Initializing database...") if err := initializeDatabase(cfg); err != nil { return fmt.Errorf("database initialization failed: %w", err) } log.Println("Database initialized successfully") // Create admin user log.Println("Creating admin user...") if err := createAdminUser(cfg); err != nil { return fmt.Errorf("admin user creation failed: %w", err) } log.Printf("Admin user created: %s", cfg.Admin.Email) // Write config file log.Println("Writing configuration file...") if err := writeConfigFile(cfg); err != nil { return fmt.Errorf("config file creation failed: %w", err) } log.Println("Configuration file created") // Create installation lock file if err := createInstallLock(); err != nil { return fmt.Errorf("failed to create install lock: %w", err) } log.Println("Installation lock created") log.Println("Auto setup completed successfully!") return nil }