First commit
This commit is contained in:
shaw
294
backend/internal/setup/cli.go
Normal file
294
backend/internal/setup/cli.go
Normal file
@@ -0,0 +1,294 @@
|
||||
package setup
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"net/mail"
|
||||
"os"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/term"
|
||||
)
|
||||
|
||||
// CLI input validation functions (matching Web API validation)
|
||||
func cliValidateHostname(host string) bool {
|
||||
validHost := regexp.MustCompile(`^[a-zA-Z0-9.\-:]+$`)
|
||||
return validHost.MatchString(host) && len(host) <= 253
|
||||
}
|
||||
|
||||
func cliValidateDBName(name string) bool {
|
||||
validName := regexp.MustCompile(`^[a-zA-Z][a-zA-Z0-9_]*$`)
|
||||
return validName.MatchString(name) && len(name) <= 63
|
||||
}
|
||||
|
||||
func cliValidateUsername(name string) bool {
|
||||
validName := regexp.MustCompile(`^[a-zA-Z0-9_]+$`)
|
||||
return validName.MatchString(name) && len(name) <= 63
|
||||
}
|
||||
|
||||
func cliValidateEmail(email string) bool {
|
||||
_, err := mail.ParseAddress(email)
|
||||
return err == nil && len(email) <= 254
|
||||
}
|
||||
|
||||
func cliValidatePort(port int) bool {
|
||||
return port > 0 && port <= 65535
|
||||
}
|
||||
|
||||
func cliValidateSSLMode(mode string) bool {
|
||||
validModes := map[string]bool{
|
||||
"disable": true, "require": true, "verify-ca": true, "verify-full": true,
|
||||
}
|
||||
return validModes[mode]
|
||||
}
|
||||
|
||||
// RunCLI runs the CLI setup wizard
|
||||
func RunCLI() error {
|
||||
reader := bufio.NewReader(os.Stdin)
|
||||
|
||||
fmt.Println()
|
||||
fmt.Println("╔═══════════════════════════════════════════╗")
|
||||
fmt.Println("║ Sub2API Installation Wizard ║")
|
||||
fmt.Println("╚═══════════════════════════════════════════╝")
|
||||
fmt.Println()
|
||||
|
||||
cfg := &SetupConfig{
|
||||
Server: ServerConfig{
|
||||
Host: "0.0.0.0",
|
||||
Port: 8080,
|
||||
Mode: "release",
|
||||
},
|
||||
JWT: JWTConfig{
|
||||
ExpireHour: 24,
|
||||
},
|
||||
}
|
||||
|
||||
// Database configuration with validation
|
||||
fmt.Println("── Database Configuration ──")
|
||||
|
||||
for {
|
||||
cfg.Database.Host = promptString(reader, "PostgreSQL Host", "localhost")
|
||||
if cliValidateHostname(cfg.Database.Host) {
|
||||
break
|
||||
}
|
||||
fmt.Println(" Invalid hostname format. Use alphanumeric, dots, hyphens only.")
|
||||
}
|
||||
|
||||
for {
|
||||
cfg.Database.Port = promptInt(reader, "PostgreSQL Port", 5432)
|
||||
if cliValidatePort(cfg.Database.Port) {
|
||||
break
|
||||
}
|
||||
fmt.Println(" Invalid port. Must be between 1 and 65535.")
|
||||
}
|
||||
|
||||
for {
|
||||
cfg.Database.User = promptString(reader, "PostgreSQL User", "postgres")
|
||||
if cliValidateUsername(cfg.Database.User) {
|
||||
break
|
||||
}
|
||||
fmt.Println(" Invalid username. Use alphanumeric and underscores only.")
|
||||
}
|
||||
|
||||
cfg.Database.Password = promptPassword("PostgreSQL Password")
|
||||
|
||||
for {
|
||||
cfg.Database.DBName = promptString(reader, "Database Name", "sub2api")
|
||||
if cliValidateDBName(cfg.Database.DBName) {
|
||||
break
|
||||
}
|
||||
fmt.Println(" Invalid database name. Start with letter, use alphanumeric and underscores.")
|
||||
}
|
||||
|
||||
for {
|
||||
cfg.Database.SSLMode = promptString(reader, "SSL Mode", "disable")
|
||||
if cliValidateSSLMode(cfg.Database.SSLMode) {
|
||||
break
|
||||
}
|
||||
fmt.Println(" Invalid SSL mode. Use: disable, require, verify-ca, or verify-full.")
|
||||
}
|
||||
|
||||
fmt.Println()
|
||||
fmt.Print("Testing database connection... ")
|
||||
if err := TestDatabaseConnection(&cfg.Database); err != nil {
|
||||
fmt.Println("FAILED")
|
||||
return fmt.Errorf("database connection failed: %w", err)
|
||||
}
|
||||
fmt.Println("OK")
|
||||
|
||||
// Redis configuration with validation
|
||||
fmt.Println()
|
||||
fmt.Println("── Redis Configuration ──")
|
||||
|
||||
for {
|
||||
cfg.Redis.Host = promptString(reader, "Redis Host", "localhost")
|
||||
if cliValidateHostname(cfg.Redis.Host) {
|
||||
break
|
||||
}
|
||||
fmt.Println(" Invalid hostname format. Use alphanumeric, dots, hyphens only.")
|
||||
}
|
||||
|
||||
for {
|
||||
cfg.Redis.Port = promptInt(reader, "Redis Port", 6379)
|
||||
if cliValidatePort(cfg.Redis.Port) {
|
||||
break
|
||||
}
|
||||
fmt.Println(" Invalid port. Must be between 1 and 65535.")
|
||||
}
|
||||
|
||||
cfg.Redis.Password = promptPassword("Redis Password (optional)")
|
||||
|
||||
for {
|
||||
cfg.Redis.DB = promptInt(reader, "Redis DB", 0)
|
||||
if cfg.Redis.DB >= 0 && cfg.Redis.DB <= 15 {
|
||||
break
|
||||
}
|
||||
fmt.Println(" Invalid Redis DB. Must be between 0 and 15.")
|
||||
}
|
||||
|
||||
fmt.Println()
|
||||
fmt.Print("Testing Redis connection... ")
|
||||
if err := TestRedisConnection(&cfg.Redis); err != nil {
|
||||
fmt.Println("FAILED")
|
||||
return fmt.Errorf("redis connection failed: %w", err)
|
||||
}
|
||||
fmt.Println("OK")
|
||||
|
||||
// Admin configuration with validation
|
||||
fmt.Println()
|
||||
fmt.Println("── Admin Account ──")
|
||||
|
||||
for {
|
||||
cfg.Admin.Email = promptString(reader, "Admin Email", "admin@example.com")
|
||||
if cliValidateEmail(cfg.Admin.Email) {
|
||||
break
|
||||
}
|
||||
fmt.Println(" Invalid email format.")
|
||||
}
|
||||
|
||||
for {
|
||||
cfg.Admin.Password = promptPassword("Admin Password")
|
||||
// SECURITY: Match Web API requirement of 8 characters minimum
|
||||
if len(cfg.Admin.Password) < 8 {
|
||||
fmt.Println(" Password must be at least 8 characters")
|
||||
continue
|
||||
}
|
||||
if len(cfg.Admin.Password) > 128 {
|
||||
fmt.Println(" Password must be at most 128 characters")
|
||||
continue
|
||||
}
|
||||
confirm := promptPassword("Confirm Password")
|
||||
if cfg.Admin.Password != confirm {
|
||||
fmt.Println(" Passwords do not match")
|
||||
continue
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
// Server configuration with validation
|
||||
fmt.Println()
|
||||
fmt.Println("── Server Configuration ──")
|
||||
|
||||
for {
|
||||
cfg.Server.Port = promptInt(reader, "Server Port", 8080)
|
||||
if cliValidatePort(cfg.Server.Port) {
|
||||
break
|
||||
}
|
||||
fmt.Println(" Invalid port. Must be between 1 and 65535.")
|
||||
}
|
||||
|
||||
// Confirm and install
|
||||
fmt.Println()
|
||||
fmt.Println("── Configuration Summary ──")
|
||||
fmt.Printf("Database: %s@%s:%d/%s\n", cfg.Database.User, cfg.Database.Host, cfg.Database.Port, cfg.Database.DBName)
|
||||
fmt.Printf("Redis: %s:%d\n", cfg.Redis.Host, cfg.Redis.Port)
|
||||
fmt.Printf("Admin: %s\n", cfg.Admin.Email)
|
||||
fmt.Printf("Server: :%d\n", cfg.Server.Port)
|
||||
fmt.Println()
|
||||
|
||||
if !promptConfirm(reader, "Proceed with installation?") {
|
||||
fmt.Println("Installation cancelled")
|
||||
return nil
|
||||
}
|
||||
|
||||
fmt.Println()
|
||||
fmt.Print("Installing... ")
|
||||
if err := Install(cfg); err != nil {
|
||||
fmt.Println("FAILED")
|
||||
return err
|
||||
}
|
||||
fmt.Println("OK")
|
||||
|
||||
fmt.Println()
|
||||
fmt.Println("╔═══════════════════════════════════════════╗")
|
||||
fmt.Println("║ Installation Complete! ║")
|
||||
fmt.Println("╚═══════════════════════════════════════════╝")
|
||||
fmt.Println()
|
||||
fmt.Println("Start the server with:")
|
||||
fmt.Println(" ./sub2api")
|
||||
fmt.Println()
|
||||
fmt.Printf("Admin panel: http://localhost:%d\n", cfg.Server.Port)
|
||||
fmt.Println()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func promptString(reader *bufio.Reader, prompt, defaultVal string) string {
|
||||
if defaultVal != "" {
|
||||
fmt.Printf(" %s [%s]: ", prompt, defaultVal)
|
||||
} else {
|
||||
fmt.Printf(" %s: ", prompt)
|
||||
}
|
||||
|
||||
input, _ := reader.ReadString('\n')
|
||||
input = strings.TrimSpace(input)
|
||||
|
||||
if input == "" {
|
||||
return defaultVal
|
||||
}
|
||||
return input
|
||||
}
|
||||
|
||||
func promptInt(reader *bufio.Reader, prompt string, defaultVal int) int {
|
||||
fmt.Printf(" %s [%d]: ", prompt, defaultVal)
|
||||
|
||||
input, _ := reader.ReadString('\n')
|
||||
input = strings.TrimSpace(input)
|
||||
|
||||
if input == "" {
|
||||
return defaultVal
|
||||
}
|
||||
|
||||
val, err := strconv.Atoi(input)
|
||||
if err != nil {
|
||||
return defaultVal
|
||||
}
|
||||
return val
|
||||
}
|
||||
|
||||
func promptPassword(prompt string) string {
|
||||
fmt.Printf(" %s: ", prompt)
|
||||
|
||||
// Try to read password without echo
|
||||
if term.IsTerminal(int(os.Stdin.Fd())) {
|
||||
password, err := term.ReadPassword(int(os.Stdin.Fd()))
|
||||
fmt.Println()
|
||||
if err == nil {
|
||||
return string(password)
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to regular input
|
||||
reader := bufio.NewReader(os.Stdin)
|
||||
input, _ := reader.ReadString('\n')
|
||||
return strings.TrimSpace(input)
|
||||
}
|
||||
|
||||
func promptConfirm(reader *bufio.Reader, prompt string) bool {
|
||||
fmt.Printf("%s [y/N]: ", prompt)
|
||||
input, _ := reader.ReadString('\n')
|
||||
input = strings.TrimSpace(strings.ToLower(input))
|
||||
return input == "y" || input == "yes"
|
||||
}
|
||||
344
backend/internal/setup/handler.go
Normal file
344
backend/internal/setup/handler.go
Normal file
@@ -0,0 +1,344 @@
|
||||
package setup
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/mail"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"sub2api/internal/pkg/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// installMutex prevents concurrent installation attempts (TOCTOU protection)
|
||||
var installMutex sync.Mutex
|
||||
|
||||
// RegisterRoutes registers setup wizard routes
|
||||
func RegisterRoutes(r *gin.Engine) {
|
||||
setup := r.Group("/setup")
|
||||
{
|
||||
// Status endpoint is always accessible (read-only)
|
||||
setup.GET("/status", getStatus)
|
||||
|
||||
// All modification endpoints are protected by setupGuard
|
||||
protected := setup.Group("")
|
||||
protected.Use(setupGuard())
|
||||
{
|
||||
protected.POST("/test-db", testDatabase)
|
||||
protected.POST("/test-redis", testRedis)
|
||||
protected.POST("/install", install)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// SetupStatus represents the current setup state
|
||||
type SetupStatus struct {
|
||||
NeedsSetup bool `json:"needs_setup"`
|
||||
Step string `json:"step"`
|
||||
}
|
||||
|
||||
// getStatus returns the current setup status
|
||||
func getStatus(c *gin.Context) {
|
||||
response.Success(c, SetupStatus{
|
||||
NeedsSetup: NeedsSetup(),
|
||||
Step: "welcome",
|
||||
})
|
||||
}
|
||||
|
||||
// setupGuard middleware ensures setup endpoints are only accessible during setup mode
|
||||
func setupGuard() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if !NeedsSetup() {
|
||||
response.Error(c, http.StatusForbidden, "Setup is not allowed: system is already installed")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// validateHostname checks if a hostname/IP is safe (no injection characters)
|
||||
func validateHostname(host string) bool {
|
||||
// Allow only alphanumeric, dots, hyphens, and colons (for IPv6)
|
||||
validHost := regexp.MustCompile(`^[a-zA-Z0-9.\-:]+$`)
|
||||
return validHost.MatchString(host) && len(host) <= 253
|
||||
}
|
||||
|
||||
// validateDBName checks if database name is safe
|
||||
func validateDBName(name string) bool {
|
||||
// Allow only alphanumeric and underscores, starting with letter
|
||||
validName := regexp.MustCompile(`^[a-zA-Z][a-zA-Z0-9_]*$`)
|
||||
return validName.MatchString(name) && len(name) <= 63
|
||||
}
|
||||
|
||||
// validateUsername checks if username is safe
|
||||
func validateUsername(name string) bool {
|
||||
// Allow only alphanumeric and underscores
|
||||
validName := regexp.MustCompile(`^[a-zA-Z0-9_]+$`)
|
||||
return validName.MatchString(name) && len(name) <= 63
|
||||
}
|
||||
|
||||
// validateEmail checks if email format is valid
|
||||
func validateEmail(email string) bool {
|
||||
_, err := mail.ParseAddress(email)
|
||||
return err == nil && len(email) <= 254
|
||||
}
|
||||
|
||||
// validatePassword checks password strength
|
||||
func validatePassword(password string) error {
|
||||
if len(password) < 8 {
|
||||
return fmt.Errorf("password must be at least 8 characters")
|
||||
}
|
||||
if len(password) > 128 {
|
||||
return fmt.Errorf("password must be at most 128 characters")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// validatePort checks if port is in valid range
|
||||
func validatePort(port int) bool {
|
||||
return port > 0 && port <= 65535
|
||||
}
|
||||
|
||||
// validateSSLMode checks if SSL mode is valid
|
||||
func validateSSLMode(mode string) bool {
|
||||
validModes := map[string]bool{
|
||||
"disable": true, "require": true, "verify-ca": true, "verify-full": true,
|
||||
}
|
||||
return validModes[mode]
|
||||
}
|
||||
|
||||
// TestDatabaseRequest represents database test request
|
||||
type TestDatabaseRequest struct {
|
||||
Host string `json:"host" binding:"required"`
|
||||
Port int `json:"port" binding:"required"`
|
||||
User string `json:"user" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
DBName string `json:"dbname" binding:"required"`
|
||||
SSLMode string `json:"sslmode"`
|
||||
}
|
||||
|
||||
// testDatabase tests database connection
|
||||
func testDatabase(c *gin.Context) {
|
||||
var req TestDatabaseRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Security: Validate all inputs to prevent injection attacks
|
||||
if !validateHostname(req.Host) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid hostname format")
|
||||
return
|
||||
}
|
||||
if !validatePort(req.Port) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid port number")
|
||||
return
|
||||
}
|
||||
if !validateUsername(req.User) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid username format")
|
||||
return
|
||||
}
|
||||
if !validateDBName(req.DBName) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid database name format")
|
||||
return
|
||||
}
|
||||
|
||||
if req.SSLMode == "" {
|
||||
req.SSLMode = "disable"
|
||||
}
|
||||
if !validateSSLMode(req.SSLMode) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid SSL mode")
|
||||
return
|
||||
}
|
||||
|
||||
cfg := &DatabaseConfig{
|
||||
Host: req.Host,
|
||||
Port: req.Port,
|
||||
User: req.User,
|
||||
Password: req.Password,
|
||||
DBName: req.DBName,
|
||||
SSLMode: req.SSLMode,
|
||||
}
|
||||
|
||||
if err := TestDatabaseConnection(cfg); err != nil {
|
||||
response.Error(c, http.StatusBadRequest, "Connection failed: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, gin.H{"message": "Connection successful"})
|
||||
}
|
||||
|
||||
// TestRedisRequest represents Redis test request
|
||||
type TestRedisRequest struct {
|
||||
Host string `json:"host" binding:"required"`
|
||||
Port int `json:"port" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
DB int `json:"db"`
|
||||
}
|
||||
|
||||
// testRedis tests Redis connection
|
||||
func testRedis(c *gin.Context) {
|
||||
var req TestRedisRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Security: Validate inputs
|
||||
if !validateHostname(req.Host) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid hostname format")
|
||||
return
|
||||
}
|
||||
if !validatePort(req.Port) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid port number")
|
||||
return
|
||||
}
|
||||
if req.DB < 0 || req.DB > 15 {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid Redis database number (0-15)")
|
||||
return
|
||||
}
|
||||
|
||||
cfg := &RedisConfig{
|
||||
Host: req.Host,
|
||||
Port: req.Port,
|
||||
Password: req.Password,
|
||||
DB: req.DB,
|
||||
}
|
||||
|
||||
if err := TestRedisConnection(cfg); err != nil {
|
||||
response.Error(c, http.StatusBadRequest, "Connection failed: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, gin.H{"message": "Connection successful"})
|
||||
}
|
||||
|
||||
// InstallRequest represents installation request
|
||||
type InstallRequest struct {
|
||||
Database DatabaseConfig `json:"database" binding:"required"`
|
||||
Redis RedisConfig `json:"redis" binding:"required"`
|
||||
Admin AdminConfig `json:"admin" binding:"required"`
|
||||
Server ServerConfig `json:"server"`
|
||||
}
|
||||
|
||||
// install performs the installation
|
||||
func install(c *gin.Context) {
|
||||
// TOCTOU Protection: Acquire mutex to prevent concurrent installation
|
||||
installMutex.Lock()
|
||||
defer installMutex.Unlock()
|
||||
|
||||
// Double-check after acquiring lock
|
||||
if !NeedsSetup() {
|
||||
response.Error(c, http.StatusForbidden, "Setup is not allowed: system is already installed")
|
||||
return
|
||||
}
|
||||
|
||||
var req InstallRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// ========== COMPREHENSIVE INPUT VALIDATION ==========
|
||||
// Database validation
|
||||
if !validateHostname(req.Database.Host) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid database hostname")
|
||||
return
|
||||
}
|
||||
if !validatePort(req.Database.Port) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid database port")
|
||||
return
|
||||
}
|
||||
if !validateUsername(req.Database.User) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid database username")
|
||||
return
|
||||
}
|
||||
if !validateDBName(req.Database.DBName) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid database name")
|
||||
return
|
||||
}
|
||||
|
||||
// Redis validation
|
||||
if !validateHostname(req.Redis.Host) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid Redis hostname")
|
||||
return
|
||||
}
|
||||
if !validatePort(req.Redis.Port) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid Redis port")
|
||||
return
|
||||
}
|
||||
if req.Redis.DB < 0 || req.Redis.DB > 15 {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid Redis database number")
|
||||
return
|
||||
}
|
||||
|
||||
// Admin validation
|
||||
if !validateEmail(req.Admin.Email) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid admin email format")
|
||||
return
|
||||
}
|
||||
if err := validatePassword(req.Admin.Password); err != nil {
|
||||
response.Error(c, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Server validation
|
||||
if req.Server.Port != 0 && !validatePort(req.Server.Port) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid server port")
|
||||
return
|
||||
}
|
||||
|
||||
// ========== SET DEFAULTS ==========
|
||||
if req.Database.SSLMode == "" {
|
||||
req.Database.SSLMode = "disable"
|
||||
}
|
||||
if !validateSSLMode(req.Database.SSLMode) {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid SSL mode")
|
||||
return
|
||||
}
|
||||
if req.Server.Host == "" {
|
||||
req.Server.Host = "0.0.0.0"
|
||||
}
|
||||
if req.Server.Port == 0 {
|
||||
req.Server.Port = 8080
|
||||
}
|
||||
if req.Server.Mode == "" {
|
||||
req.Server.Mode = "release"
|
||||
}
|
||||
// Validate server mode
|
||||
if req.Server.Mode != "release" && req.Server.Mode != "debug" {
|
||||
response.Error(c, http.StatusBadRequest, "Invalid server mode (must be 'release' or 'debug')")
|
||||
return
|
||||
}
|
||||
|
||||
// Trim whitespace from string inputs
|
||||
req.Admin.Email = strings.TrimSpace(req.Admin.Email)
|
||||
req.Database.Host = strings.TrimSpace(req.Database.Host)
|
||||
req.Database.User = strings.TrimSpace(req.Database.User)
|
||||
req.Database.DBName = strings.TrimSpace(req.Database.DBName)
|
||||
req.Redis.Host = strings.TrimSpace(req.Redis.Host)
|
||||
|
||||
cfg := &SetupConfig{
|
||||
Database: req.Database,
|
||||
Redis: req.Redis,
|
||||
Admin: req.Admin,
|
||||
Server: req.Server,
|
||||
JWT: JWTConfig{
|
||||
ExpireHour: 24,
|
||||
},
|
||||
}
|
||||
|
||||
if err := Install(cfg); err != nil {
|
||||
response.Error(c, http.StatusInternalServerError, "Installation failed: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, gin.H{
|
||||
"message": "Installation completed successfully",
|
||||
"restart": true,
|
||||
})
|
||||
}
|
||||
564
backend/internal/setup/setup.go
Normal file
564
backend/internal/setup/setup.go
Normal file
@@ -0,0 +1,564 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user