option backup restore #238

2025-01-20 01:26:02 +01:00
parent 049cfc5287
commit f3432b119c
11 changed files with 442 additions and 33 deletions
@@ -0,0 +1,271 @@
+package database
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"mime/multipart"
+	"os"
+	"path/filepath"
+	"s-ui/cmd/migration"
+	"s-ui/config"
+	"s-ui/database/model"
+	"s-ui/logger"
+	"s-ui/util/common"
+	"strings"
+	"syscall"
+	"time"
+
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+)
+
+func GetDb(exclude string) ([]byte, error) {
+	exclude_changes, exclude_stats := false, false
+	for _, table := range strings.Split(exclude, ",") {
+		if table == "changes" {
+			exclude_changes = true
+		} else if table == "stats" {
+			exclude_stats = true
+		}
+	}
+
+	dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
+	if err != nil {
+		return nil, err
+	}
+	dbPath := dir + config.GetName() + "_" + time.Now().Format("20060102-200203") + ".db"
+
+	backupDb, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
+	if err != nil {
+		return nil, err
+	}
+
+	err = backupDb.AutoMigrate(
+		&model.Setting{},
+		&model.Tls{},
+		&model.Inbound{},
+		&model.Outbound{},
+		&model.Endpoint{},
+		&model.User{},
+		&model.Stats{},
+		&model.Client{},
+		&model.Changes{},
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	var settings []model.Setting
+	var tls []model.Tls
+	var inbound []model.Inbound
+	var outbound []model.Outbound
+	var endpoint []model.Endpoint
+	var users []model.User
+	var clients []model.Client
+	var stats []model.Stats
+	var changes []model.Changes
+
+	// Perform scans and handle errors
+	if err := db.Model(&model.Setting{}).Scan(&settings).Error; err != nil {
+		return nil, err
+	}
+	if err := db.Model(&model.Tls{}).Scan(&tls).Error; err != nil {
+		return nil, err
+	}
+	if err := db.Model(&model.Inbound{}).Scan(&inbound).Error; err != nil {
+		return nil, err
+	}
+	if err := db.Model(&model.Outbound{}).Scan(&outbound).Error; err != nil {
+		return nil, err
+	}
+	if err := db.Model(&model.Endpoint{}).Scan(&endpoint).Error; err != nil {
+		return nil, err
+	}
+	if err := db.Model(&model.User{}).Scan(&users).Error; err != nil {
+		return nil, err
+	}
+	if err := db.Model(&model.Client{}).Scan(&clients).Error; err != nil {
+		return nil, err
+	}
+
+	// Save each model
+	for _, mdl := range []interface{}{settings, tls, inbound, outbound, endpoint, users, clients} {
+		if err := backupDb.Save(mdl).Error; err != nil {
+			return nil, err
+		}
+	}
+
+	if !exclude_stats {
+		if err := db.Model(&model.Stats{}).Scan(&stats).Error; err != nil {
+			return nil, err
+		}
+		if err := backupDb.Save(stats).Error; err != nil {
+			return nil, err
+		}
+	}
+	if !exclude_changes {
+		if err := db.Model(&model.Changes{}).Scan(&changes).Error; err != nil {
+			return nil, err
+		}
+		if err := backupDb.Save(changes).Error; err != nil {
+			return nil, err
+		}
+	}
+
+	// Update WAL
+	err = backupDb.Exec("PRAGMA wal_checkpoint;").Error
+	if err != nil {
+		return nil, err
+	}
+
+	bdb, _ := backupDb.DB()
+	bdb.Close()
+
+	// Open the file for reading
+	file, err := os.Open(dbPath)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+	defer os.Remove(dbPath)
+
+	// Read the file contents
+	fileContents, err := io.ReadAll(file)
+	if err != nil {
+		return nil, err
+	}
+
+	return fileContents, nil
+}
+
+func ImportDB(file multipart.File) error {
+	// Check if the file is a SQLite database
+	isValidDb, err := IsSQLiteDB(file)
+	if err != nil {
+		return common.NewErrorf("Error checking db file format: %v", err)
+	}
+	if !isValidDb {
+		return common.NewError("Invalid db file format")
+	}
+
+	// Reset the file reader to the beginning
+	_, err = file.Seek(0, 0)
+	if err != nil {
+		return common.NewErrorf("Error resetting file reader: %v", err)
+	}
+
+	// Save the file as temporary file
+	tempPath := fmt.Sprintf("%s.temp", config.GetDBPath())
+	// Remove the existing fallback file (if any) before creating one
+	_, err = os.Stat(tempPath)
+	if err == nil {
+		errRemove := os.Remove(tempPath)
+		if errRemove != nil {
+			return common.NewErrorf("Error removing existing temporary db file: %v", errRemove)
+		}
+	}
+	// Create the temporary file
+	tempFile, err := os.Create(tempPath)
+	if err != nil {
+		return common.NewErrorf("Error creating temporary db file: %v", err)
+	}
+	defer tempFile.Close()
+
+	// Remove temp file before returning
+	defer os.Remove(tempPath)
+
+	// Close old DB
+	old_db, _ := db.DB()
+	old_db.Close()
+
+	// Save uploaded file to temporary file
+	_, err = io.Copy(tempFile, file)
+	if err != nil {
+		return common.NewErrorf("Error saving db: %v", err)
+	}
+
+	// Check if we can init db or not
+	newDb, err := gorm.Open(sqlite.Open(tempPath), &gorm.Config{})
+	if err != nil {
+		return common.NewErrorf("Error checking db: %v", err)
+	}
+	newDb_db, _ := newDb.DB()
+	newDb_db.Close()
+
+	// Backup the current database for fallback
+	fallbackPath := fmt.Sprintf("%s.backup", config.GetDBPath())
+	// Remove the existing fallback file (if any)
+	_, err = os.Stat(fallbackPath)
+	if err == nil {
+		errRemove := os.Remove(fallbackPath)
+		if errRemove != nil {
+			return common.NewErrorf("Error removing existing fallback db file: %v", errRemove)
+		}
+	}
+	// Move the current database to the fallback location
+	err = os.Rename(config.GetDBPath(), fallbackPath)
+	if err != nil {
+		return common.NewErrorf("Error backing up temporary db file: %v", err)
+	}
+
+	// Remove the temporary file before returning
+	defer os.Remove(fallbackPath)
+
+	// Move temp to DB path
+	err = os.Rename(tempPath, config.GetDBPath())
+	if err != nil {
+		errRename := os.Rename(fallbackPath, config.GetDBPath())
+		if errRename != nil {
+			return common.NewErrorf("Error moving db file and restoring fallback: %v", errRename)
+		}
+		return common.NewErrorf("Error moving db file: %v", err)
+	}
+
+	// Migrate DB
+	migration.MigrateDb()
+	err = InitDB(config.GetDBPath())
+	if err != nil {
+		errRename := os.Rename(fallbackPath, config.GetDBPath())
+		if errRename != nil {
+			return common.NewErrorf("Error migrating db and restoring fallback: %v", errRename)
+		}
+		return common.NewErrorf("Error migrating db: %v", err)
+	}
+
+	// Restart app
+	err = SendSighup()
+	if err != nil {
+		return common.NewErrorf("Error restarting app: %v", err)
+	}
+
+	return nil
+}
+
+func IsSQLiteDB(file io.Reader) (bool, error) {
+	signature := []byte("SQLite format 3\x00")
+	buf := make([]byte, len(signature))
+	_, err := file.Read(buf)
+	if err != nil {
+		return false, err
+	}
+	return bytes.Equal(buf, signature), nil
+}
+
+func SendSighup() error {
+	// Get the current process
+	process, err := os.FindProcess(os.Getpid())
+	if err != nil {
+		return err
+	}
+
+	// Send SIGHUP to the current process
+	go func() {
+		time.Sleep(3 * time.Second)
+		err := process.Signal(syscall.SIGHUP)
+		if err != nil {
+			logger.Error("send signal SIGHUP failed:", err)
+		}
+	}()
+	return nil
+}