TaskPool/internal/services/backup_service.go

package services

import (
	"archive/zip"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"

	"baihu/internal/constant"
	"baihu/internal/database"
	"baihu/internal/models"
)

type BackupService struct {
	settingsService *SettingsService
}

func NewBackupService() *BackupService {
	return &BackupService{
		settingsService: NewSettingsService(),
	}
}

const (
	BackupSection = "backup"
	BackupFileKey = "backup_file"
	BackupDir     = "./data/backups"
)

// tableConfig 表备份配置
type tableConfig struct {
	filename string
	export   func() (any, error)
	restore  func([]byte) error
}

func (s *BackupService) getTableConfigs() []tableConfig {
	return []tableConfig{
		{"tasks.json", s.exportTable(&[]models.Task{}, true), s.restoreTable(&[]models.Task{}, true)},
		{"task_logs.json", s.exportTable(&[]models.TaskLog{}, false), s.restoreTable(&[]models.TaskLog{}, false)},
		{"envs.json", s.exportTable(&[]models.EnvironmentVariable{}, true), s.restoreTable(&[]models.EnvironmentVariable{}, true)},
		{"scripts.json", s.exportTable(&[]models.Script{}, true), s.restoreTable(&[]models.Script{}, true)},
		{"settings.json", s.exportSettings, s.restoreSettings},
		{"send_stats.json", s.exportTable(&[]models.SendStats{}, false), s.restoreTable(&[]models.SendStats{}, false)},
		{"login_logs.json", s.exportTable(&[]models.LoginLog{}, false), s.restoreTable(&[]models.LoginLog{}, false)},
	}
}

func (s *BackupService) exportTable(dest any, unscoped bool) func() (any, error) {
	return func() (any, error) {
		db := database.DB
		if unscoped {
			db = db.Unscoped()
		}
		db.Find(dest)
		return dest, nil
	}
}

func (s *BackupService) restoreTable(dest any, unscoped bool) func([]byte) error {
	return func(data []byte) error {
		if err := json.Unmarshal(data, dest); err != nil {
			return err
		}
		return nil
	}
}

func (s *BackupService) exportSettings() (any, error) {
	var data []models.Setting
	database.DB.Where("section != ?", BackupSection).Find(&data)
	return data, nil
}

func (s *BackupService) restoreSettings(data []byte) error {
	var settings []models.Setting
	return json.Unmarshal(data, &settings)
}

// CreateBackup 创建备份
func (s *BackupService) CreateBackup() (string, error) {
	if err := os.MkdirAll(BackupDir, 0755); err != nil {
		return "", err
	}

	timestamp := time.Now().Format("20060102_150405")
	zipPath := filepath.Join(BackupDir, fmt.Sprintf("backup_%s.zip", timestamp))

	zipFile, err := os.Create(zipPath)
	if err != nil {
		return "", err
	}
	defer zipFile.Close()

	zipWriter := zip.NewWriter(zipFile)
	defer zipWriter.Close()

	// 导出各表
	for _, cfg := range s.getTableConfigs() {
		data, err := cfg.export()
		if err != nil {
			return "", err
		}
		jsonData, err := json.MarshalIndent(data, "", "  ")
		if err != nil {
			return "", err
		}
		w, err := zipWriter.Create(cfg.filename)
		if err != nil {
			return "", err
		}
		if _, err := w.Write(jsonData); err != nil {
			return "", err
		}
	}

	// 打包 scripts 文件夹
	scriptsDir := constant.ScriptsWorkDir
	if _, err := os.Stat(scriptsDir); err == nil {
		if err := s.addDirToZip(zipWriter, scriptsDir, "scripts"); err != nil {
			return "", err
		}
	}

	s.settingsService.Set(BackupSection, BackupFileKey, zipPath)
	return zipPath, nil
}

// Restore 恢复备份
func (s *BackupService) Restore(zipPath string) error {
	r, err := zip.OpenReader(zipPath)
	if err != nil {
		return err
	}
	defer r.Close()

	// 构建文件名到配置的映射
	configs := s.getTableConfigs()
	fileMap := make(map[string]*zip.File)
	for _, f := range r.File {
		fileMap[f.Name] = f
	}

	// 读取所有表数据
	tableData := make(map[string][]byte)
	for _, cfg := range configs {
		if f, ok := fileMap[cfg.filename]; ok {
			data, err := s.readZipFile(f)
			if err != nil {
				return err
			}
			tableData[cfg.filename] = data
		}
	}

	// 清空现有数据（物理删除）
	database.DB.Unscoped().Where("1=1").Delete(&models.Task{})
	database.DB.Unscoped().Where("1=1").Delete(&models.TaskLog{})
	database.DB.Unscoped().Where("1=1").Delete(&models.EnvironmentVariable{})
	database.DB.Unscoped().Where("1=1").Delete(&models.Script{})
	database.DB.Unscoped().Where("section != ?", BackupSection).Delete(&models.Setting{})
	database.DB.Unscoped().Where("1=1").Delete(&models.SendStats{})
	database.DB.Unscoped().Where("1=1").Delete(&models.LoginLog{})

	// 恢复数据
	s.restoreFromData(tableData, "tasks.json", &[]models.Task{})
	s.restoreFromData(tableData, "task_logs.json", &[]models.TaskLog{})
	s.restoreFromData(tableData, "envs.json", &[]models.EnvironmentVariable{})
	s.restoreFromData(tableData, "scripts.json", &[]models.Script{})
	s.restoreFromData(tableData, "settings.json", &[]models.Setting{})
	s.restoreFromData(tableData, "send_stats.json", &[]models.SendStats{})
	s.restoreFromData(tableData, "login_logs.json", &[]models.LoginLog{})

	// 恢复 scripts 文件夹
	s.restoreScriptsDir(r)

	return nil
}

func (s *BackupService) restoreFromData(tableData map[string][]byte, filename string, dest any) {
	if data, ok := tableData[filename]; ok {
		if err := json.Unmarshal(data, dest); err == nil {
			s.insertRecords(dest)
		}
	}
}

func (s *BackupService) insertRecords(records any) {
	switch v := records.(type) {
	case *[]models.Task:
		for _, r := range *v {
			database.DB.Create(&r)
		}
	case *[]models.TaskLog:
		for _, r := range *v {
			database.DB.Create(&r)
		}
	case *[]models.EnvironmentVariable:
		for _, r := range *v {
			database.DB.Create(&r)
		}
	case *[]models.Script:
		for _, r := range *v {
			database.DB.Create(&r)
		}
	case *[]models.Setting:
		for _, r := range *v {
			database.DB.Create(&r)
		}
	case *[]models.SendStats:
		for _, r := range *v {
			database.DB.Create(&r)
		}
	case *[]models.LoginLog:
		for _, r := range *v {
			database.DB.Create(&r)
		}
	}
}

func (s *BackupService) restoreScriptsDir(r *zip.ReadCloser) {
	scriptsDir := constant.ScriptsWorkDir
	for _, f := range r.File {
		if len(f.Name) > 8 && f.Name[:8] == "scripts/" {
			relPath := f.Name[8:]
			if relPath == "" {
				continue
			}
			fpath := filepath.Join(scriptsDir, relPath)
			if f.FileInfo().IsDir() {
				os.MkdirAll(fpath, 0755)
				continue
			}
			os.MkdirAll(filepath.Dir(fpath), 0755)
			if outFile, err := os.Create(fpath); err == nil {
				if rc, err := f.Open(); err == nil {
					io.Copy(outFile, rc)
					rc.Close()
				}
				outFile.Close()
			}
		}
	}
}

func (s *BackupService) readZipFile(f *zip.File) ([]byte, error) {
	rc, err := f.Open()
	if err != nil {
		return nil, err
	}
	defer rc.Close()
	return io.ReadAll(rc)
}

func (s *BackupService) addDirToZip(zipWriter *zip.Writer, srcDir, prefix string) error {
	return filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		relPath, err := filepath.Rel(srcDir, path)
		if err != nil {
			return err
		}
		zipPath := filepath.ToSlash(filepath.Join(prefix, relPath))
		if info.IsDir() {
			if relPath != "." {
				_, err := zipWriter.Create(zipPath + "/")
				return err
			}
			return nil
		}
		w, err := zipWriter.Create(zipPath)
		if err != nil {
			return err
		}
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()
		_, err = io.Copy(w, file)
		return err
	})
}

func (s *BackupService) GetBackupFile() string {
	var setting models.Setting
	if err := database.DB.Where("section = ? AND `key` = ?", BackupSection, BackupFileKey).First(&setting).Error; err != nil {
		return ""
	}
	return setting.Value
}

func (s *BackupService) ClearBackup() error {
	filePath := s.GetBackupFile()
	if filePath != "" {
		os.Remove(filePath)
		database.DB.Where("section = ? AND `key` = ?", BackupSection, BackupFileKey).Delete(&models.Setting{})
	}
	return nil
}