LokiStack/backend/services/backup.go

package services

import (
	"archive/tar"
	"compress/gzip"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"lokistack/repository"
)

// 备份/恢复设计要点
//
// 为什么必须有：平台的全部状态（部署/采集/存储/集群/主机/策略/告警/用户/审计）
// 都在这一个 JSON 文件里；kubeconfig 虽然加密存在 DB 内，但解密依赖
// ENCRYPTION_KEY（K8s Secret）。DB 丢了 = 平台失去"这些资源是谁建的"的映射，
// 从此无法管理/卸载/升级正在跑的 Loki 与采集器。
//
// 安全取舍：
//   - 备份包内含加密后的 kubeconfig 与 known_hosts，**不含明文密钥**
//   - manifest 里只记录 ENCRYPTION_KEY 的指纹（sha256 前 16 位），
//     用于恢复时校验"这把钥匙能不能解开这份备份"，不泄露密钥本身
//   - 恢复分两步（校验暂存 → 手动应用），避免一个误点就覆盖线上 DB

// BackupManifest 备份包元信息（放在包内 manifest.json）
type BackupManifest struct {
	Version     string    `json:"version"`
	CreatedAt   time.Time `json:"createdAt"`
	DBBytes     int64     `json:"dbBytes"`
	DBSHA256    string    `json:"dbSha256"`
	KeyFP       string    `json:"encryptionKeyFingerprint"`
	RecordCount int       `json:"recordCount"`
	Note        string    `json:"note"`
}

// backupExclude 备份时跳过的文件/目录（临时文件、lost+found、正在写的 restore 暂存）
var backupExclude = map[string]bool{
	"lost+found": true, ".restore-pending": true, "backups": true,
}

// CreateBackup 把 /app/data 打包成 tar.gz 写入 w。
func CreateBackup(w io.Writer) (*BackupManifest, error) {
	dataDir := dataDirPath()
	dbPath := databasePath()

	dbInfo, err := os.Stat(dbPath)
	if err != nil {
		return nil, fmt.Errorf("读取数据库失败: %w", err)
	}
	dbBytes, err := os.ReadFile(dbPath)
	if err != nil {
		return nil, fmt.Errorf("读取数据库失败: %w", err)
	}
	sum := sha256.Sum256(dbBytes)

	// 记录数从**备份出来的字节**里数，不用 repository.GetDB()：
	// CLI 模式（lokistack backup，跑在 CronJob 里）根本没初始化 DB，
	// 用 GetDB() 会得到 nil，manifest 里的 recordCount 恒为 0，
	// 恢复时就看不出"这份备份里有没有东西"。
	recordCount := countRecords(dbBytes)

	manifest := &BackupManifest{
		Version:     "2.0.0",
		CreatedAt:   time.Now(),
		DBBytes:     dbInfo.Size(),
		DBSHA256:    hex.EncodeToString(sum[:]),
		KeyFP:       encryptionKeyFingerprint(),
		RecordCount: recordCount,
		Note:        "包内 kubeconfig 为加密态；恢复必须使用同一把 ENCRYPTION_KEY，否则 kubeconfig 无法解密",
	}

	gz := gzip.NewWriter(w)
	defer gz.Close()
	tw := tar.NewWriter(gz)
	defer tw.Close()

	// 1) manifest
	mBytes, _ := json.MarshalIndent(manifest, "", "  ")
	if err := writeTarFile(tw, "manifest.json", mBytes, 0o600); err != nil {
		return nil, err
	}

	// 2) data 目录下所有文件（DB + kubeconfigs/ + known_hosts/）
	err = filepath.Walk(dataDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil // 单个文件读不到不阻断整包
		}
		if path == dataDir {
			return nil
		}
		rel, err := filepath.Rel(dataDir, path)
		if err != nil {
			return nil
		}
		top := strings.SplitN(rel, string(os.PathSeparator), 2)[0]
		if backupExclude[top] {
			if info.IsDir() {
				return filepath.SkipDir
			}
			return nil
		}
		if info.IsDir() {
			return nil
		}
		// 跳过正在写的临时文件
		if strings.HasPrefix(filepath.Base(path), ".lokistack-") && strings.HasSuffix(path, ".tmp") {
			return nil
		}
		f, err := os.Open(path)
		if err != nil {
			return nil
		}
		defer f.Close()
		hdr := &tar.Header{
			Name:    filepath.ToSlash(filepath.Join("data", rel)),
			Mode:    int64(info.Mode().Perm()),
			Size:    info.Size(),
			ModTime: info.ModTime(),
		}
		if err := tw.WriteHeader(hdr); err != nil {
			return err
		}
		_, err = io.Copy(tw, f)
		return err
	})
	if err != nil {
		return nil, err
	}
	return manifest, nil
}

func writeTarFile(tw *tar.Writer, name string, data []byte, mode int64) error {
	hdr := &tar.Header{Name: name, Mode: mode, Size: int64(len(data)), ModTime: time.Now()}
	if err := tw.WriteHeader(hdr); err != nil {
		return err
	}
	_, err := tw.Write(data)
	return err
}

// RestoreResult 校验结果
type RestoreResult struct {
	Manifest     *BackupManifest `json:"manifest"`
	KeyMatches   bool            `json:"keyMatches"`
	RowsInBackup int             `json:"rowsInBackup"`
	CurrentRows  int             `json:"currentRows"`
	StagedPath   string          `json:"stagedPath,omitempty"`
	Warning      string          `json:"warning,omitempty"`
}

// ValidateAndStageBackup 校验备份包并暂存到 <dataDir>/.restore-pending/。
//
// 不在这一步直接覆盖线上 DB：恢复是不可逆操作，必须让人再确认一次
// （页面上的「应用恢复」或 run.sh restore）。
func ValidateAndStageBackup(r io.Reader) (*RestoreResult, error) {
	gr, err := gzip.NewReader(r)
	if err != nil {
		return nil, fmt.Errorf("不是合法的 gzip 文件: %w", err)
	}
	defer gr.Close()
	tr := tar.NewReader(gr)

	var (
		manifest *BackupManifest
		dbBytes  []byte
	)
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, fmt.Errorf("解包失败: %w", err)
		}
		switch {
		case hdr.Name == "manifest.json":
			b, _ := io.ReadAll(io.LimitReader(tr, 1<<20))
			var m BackupManifest
			if err := json.Unmarshal(b, &m); err == nil {
				manifest = &m
			}
		case hdr.Name == "data/opsxloki.json":
			dbBytes, _ = io.ReadAll(io.LimitReader(tr, 64<<20))
		}
	}
	if manifest == nil {
		return nil, fmt.Errorf("备份包缺少 manifest.json，不是本平台导出的备份")
	}
	if len(dbBytes) == 0 {
		return nil, fmt.Errorf("备份包内没有 data/opsxloki.json")
	}

	// 校验 DB 完整性：必须是合法 JSON 且能解析出预期集合
	var probe map[string]json.RawMessage
	if err := json.Unmarshal(dbBytes, &probe); err != nil {
		return nil, fmt.Errorf("备份内的数据库不是合法 JSON: %w", err)
	}
	for _, required := range []string{"users", "deployments", "collectors"} {
		if _, ok := probe[required]; !ok {
			return nil, fmt.Errorf("备份内数据库缺少集合 %q，可能不是 LokiStack 的备份", required)
		}
	}

	// 哈希校验（防止传输损坏）
	sum := sha256.Sum256(dbBytes)
	if manifest.DBSHA256 != "" && hex.EncodeToString(sum[:]) != manifest.DBSHA256 {
		return nil, fmt.Errorf("备份内数据库校验失败（sha256 不匹配），文件可能已损坏")
	}

	stagedDir := filepath.Join(dataDirPath(), ".restore-pending")
	if err := os.MkdirAll(stagedDir, 0o700); err != nil {
		return nil, fmt.Errorf("创建暂存目录失败: %w", err)
	}
	stagedPath := filepath.Join(stagedDir, "opsxloki.json")
	if err := os.WriteFile(stagedPath, dbBytes, 0o600); err != nil {
		return nil, fmt.Errorf("写入暂存文件失败: %w", err)
	}

	currentRows := 0
	if db := repository.GetDB(); db != nil {
		db.RLock()
		currentRows = len(db.Deployments) + len(db.Collectors) + len(db.Clusters) + len(db.Hosts)
		db.RUnlock()
	}

	res := &RestoreResult{
		Manifest:     manifest,
		KeyMatches:   manifest.KeyFP == "" || manifest.KeyFP == encryptionKeyFingerprint(),
		RowsInBackup: manifest.RecordCount,
		CurrentRows:  currentRows,
		StagedPath:   stagedPath,
	}
	if !res.KeyMatches {
		res.Warning = "备份的 ENCRYPTION_KEY 指纹与当前平台不一致：恢复后 kubeconfig / 存储凭证将无法解密，" +
			"需要用备份时的那把 ENCRYPTION_KEY 覆盖 lokistack-secrets 中的值后再启动"
	}
	return res, nil
}

// countRecords 从数据库 JSON 字节里统计各集合的记录总数。
func countRecords(dbBytes []byte) int {
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(dbBytes, &raw); err != nil {
		return 0
	}
	total := 0
	for _, v := range raw {
		var arr []json.RawMessage
		if err := json.Unmarshal(v, &arr); err == nil {
			total += len(arr)
		}
	}
	return total
}

// encryptionKeyFingerprint 返回当前 ENCRYPTION_KEY 的指纹（sha256 前 16 位），不泄露密钥。
func encryptionKeyFingerprint() string {
	key := os.Getenv("ENCRYPTION_KEY")
	if key == "" {
		// 兼容从文件读取的部署方式
		if p := os.Getenv("ENCRYPTION_KEY_FILE"); p != "" {
			if b, err := os.ReadFile(p); err == nil {
				key = strings.TrimSpace(string(b))
			}
		}
	}
	if key == "" {
		return ""
	}
	sum := sha256.Sum256([]byte(key))
	return hex.EncodeToString(sum[:])[:16]
}

// ApplyStagedBackup 把暂存文件应用为正式 DB（调用方负责重启进程）。
func ApplyStagedBackup() (string, error) {
	staged := filepath.Join(dataDirPath(), ".restore-pending", "opsxloki.json")
	b, err := os.ReadFile(staged)
	if err != nil {
		return "", fmt.Errorf("没有待应用的恢复文件: %w", err)
	}
	var probe map[string]json.RawMessage
	if err := json.Unmarshal(b, &probe); err != nil {
		return "", fmt.Errorf("暂存文件已损坏: %w", err)
	}

	dbPath := databasePath()
	// 先把当前 DB 备份一份，恢复错了还能退回去
	if cur, err := os.ReadFile(dbPath); err == nil {
		rollback := dbPath + ".before-restore." + time.Now().Format("20060102-150405")
		_ = os.WriteFile(rollback, cur, 0o600)
	}
	tmp := dbPath + ".restoring"
	if err := os.WriteFile(tmp, b, 0o600); err != nil {
		return "", err
	}
	if err := os.Rename(tmp, dbPath); err != nil {
		return "", err
	}
	_ = os.Remove(staged)
	return dbPath, nil
}

// ListBackups 列出 dataDir/backups 下已有的备份文件（run.sh backup 产出）。
func ListBackups() []map[string]interface{} {
	dir := filepath.Join(dataDirPath(), "backups")
	entries, err := os.ReadDir(dir)
	if err != nil {
		return []map[string]interface{}{}
	}
	type item struct {
		name string
		size int64
		mod  time.Time
	}
	var items []item
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(e.Name(), ".tar.gz") {
			continue
		}
		info, err := e.Info()
		if err != nil {
			continue
		}
		items = append(items, item{e.Name(), info.Size(), info.ModTime()})
	}
	sort.Slice(items, func(i, j int) bool { return items[i].mod.After(items[j].mod) })
	out := make([]map[string]interface{}, 0, len(items))
	for _, it := range items {
		out = append(out, map[string]interface{}{
			"name": it.name, "sizeBytes": it.size, "createdAt": it.mod,
		})
	}
	return out
}

func dataDirPath() string {
	if p := os.Getenv("DATA_DIR"); p != "" {
		return p
	}
	return "/app/data"
}