346 lines
7.1 KiB
Go
346 lines
7.1 KiB
Go
package cmd
|
|
|
|
import (
|
|
"archive/tar"
|
|
"compress/gzip"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"strings"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/charmbracelet/bubbles/progress"
|
|
tea "github.com/charmbracelet/bubbletea"
|
|
"github.com/charmbracelet/lipgloss"
|
|
"github.com/spf13/cobra"
|
|
)
|
|
|
|
var (
|
|
restoreSource string
|
|
restoreTarget string
|
|
)
|
|
|
|
var restoreCmd = &cobra.Command{
|
|
Use: "restore",
|
|
Short: "Restore a backup",
|
|
Long: "Restore a backup from an archive to a file, disk, or directory.",
|
|
Run: func(cmd *cobra.Command, args []string) {
|
|
model, err := initialRestoreModel(restoreSource, restoreTarget)
|
|
if err != nil {
|
|
fmt.Println("Error:", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
p := tea.NewProgram(model, tea.WithAltScreen())
|
|
if _, err := p.Run(); err != nil {
|
|
fmt.Println("Error:", err)
|
|
os.Exit(1)
|
|
}
|
|
},
|
|
}
|
|
|
|
type restoreModel struct {
|
|
progress progress.Model
|
|
source string
|
|
target string
|
|
totalBytes int64
|
|
processed int64
|
|
done bool
|
|
err error
|
|
startTime time.Time
|
|
progressCh chan int64
|
|
isDirectory bool
|
|
}
|
|
|
|
type restoreProgressMsg int64
|
|
type restoreDoneMsg struct{}
|
|
type restoreErrorMsg struct{ err error }
|
|
|
|
func initialRestoreModel(src, dst string) (*restoreModel, error) {
|
|
fileInfo, err := os.Stat(src)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
p := progress.New(
|
|
progress.WithDefaultGradient(),
|
|
progress.WithWidth(40),
|
|
)
|
|
|
|
isDirectory := strings.HasSuffix(src, ".tar.gz")
|
|
targetPath := dst
|
|
|
|
if isDirectory && targetPath == "" {
|
|
targetPath, err = os.Getwd()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get current directory: %w", err)
|
|
}
|
|
} else if !isDirectory && targetPath == "" {
|
|
return nil, fmt.Errorf("target flag is required for file or disk restoration")
|
|
}
|
|
|
|
return &restoreModel{
|
|
progress: p,
|
|
source: src,
|
|
target: targetPath,
|
|
totalBytes: fileInfo.Size(),
|
|
startTime: time.Now(),
|
|
progressCh: make(chan int64, 100),
|
|
isDirectory: isDirectory,
|
|
}, nil
|
|
}
|
|
|
|
func (m *restoreModel) Init() tea.Cmd {
|
|
return tea.Batch(
|
|
m.startRestore,
|
|
m.restoreProgressListener,
|
|
)
|
|
}
|
|
|
|
func (m *restoreModel) startRestore() tea.Msg {
|
|
go func() {
|
|
if m.isDirectory {
|
|
if err := m.restoreDirectory(); err != nil {
|
|
m.progressCh <- -1 // Error signal
|
|
return
|
|
}
|
|
m.progressCh <- -2 // Done signal
|
|
return
|
|
}
|
|
|
|
// Fallback to original file/disk restoration logic
|
|
srcFile, err := os.Open(m.source)
|
|
if err != nil {
|
|
m.progressCh <- -1
|
|
return
|
|
}
|
|
defer srcFile.Close()
|
|
|
|
progressReader := &restoreProgressReader{
|
|
reader: srcFile,
|
|
total: m.totalBytes,
|
|
progressCh: m.progressCh,
|
|
}
|
|
|
|
gzipReader, err := gzip.NewReader(progressReader)
|
|
if err != nil {
|
|
m.progressCh <- -1
|
|
return
|
|
}
|
|
defer gzipReader.Close()
|
|
|
|
var dstFile *os.File
|
|
if isBlockDevice(m.target) {
|
|
dstFile, err = os.OpenFile(m.target, os.O_WRONLY|os.O_SYNC, 0)
|
|
if err != nil {
|
|
m.progressCh <- -1
|
|
return
|
|
}
|
|
} else {
|
|
dstFile, err = os.Create(m.target)
|
|
if err != nil {
|
|
m.progressCh <- -1
|
|
return
|
|
}
|
|
}
|
|
defer dstFile.Close()
|
|
|
|
_, err = io.Copy(dstFile, gzipReader)
|
|
if err != nil {
|
|
m.progressCh <- -1
|
|
return
|
|
}
|
|
|
|
if err := dstFile.Sync(); err != nil {
|
|
m.progressCh <- -1
|
|
return
|
|
}
|
|
|
|
m.progressCh <- -2 // Сигнал завершения
|
|
}()
|
|
|
|
return nil
|
|
}
|
|
|
|
func (m *restoreModel) restoreDirectory() error {
|
|
srcFile, err := os.Open(m.source)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer srcFile.Close()
|
|
|
|
progressReader := &restoreProgressReader{
|
|
reader: srcFile,
|
|
total: m.totalBytes,
|
|
progressCh: m.progressCh,
|
|
}
|
|
|
|
gzipReader, err := gzip.NewReader(progressReader)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer gzipReader.Close()
|
|
|
|
tarReader := tar.NewReader(gzipReader)
|
|
|
|
for {
|
|
header, err := tarReader.Next()
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
targetPath := filepath.Join(m.target, header.Name)
|
|
|
|
switch header.Typeflag {
|
|
case tar.TypeDir:
|
|
if err := os.MkdirAll(targetPath, os.FileMode(header.Mode)); err != nil {
|
|
return err
|
|
}
|
|
case tar.TypeReg:
|
|
if err := os.MkdirAll(filepath.Dir(targetPath), 0755); err != nil {
|
|
return err
|
|
}
|
|
outFile, err := os.OpenFile(targetPath, os.O_CREATE|os.O_WRONLY, os.FileMode(header.Mode))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if _, err := io.Copy(outFile, tarReader); err != nil {
|
|
outFile.Close()
|
|
return err
|
|
}
|
|
outFile.Close()
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type restoreProgressReader struct {
|
|
reader io.Reader
|
|
total int64
|
|
progressCh chan int64
|
|
read int64
|
|
}
|
|
|
|
func (r *restoreProgressReader) Read(p []byte) (int, error) {
|
|
n, err := r.reader.Read(p)
|
|
if n > 0 {
|
|
r.read += int64(n)
|
|
r.progressCh <- r.read
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
func (m *restoreModel) restoreProgressListener() tea.Msg {
|
|
select {
|
|
case n := <-m.progressCh:
|
|
switch {
|
|
case n == -1: // Ошибка
|
|
return restoreErrorMsg{fmt.Errorf("restoration error")}
|
|
case n == -2: // Завершение
|
|
return restoreDoneMsg{}
|
|
default: // Прогресс
|
|
m.processed = n
|
|
return restoreProgressMsg(n)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (m *restoreModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
|
switch msg := msg.(type) {
|
|
case tea.KeyMsg:
|
|
if msg.Type == tea.KeyCtrlC {
|
|
return m, tea.Quit
|
|
}
|
|
return m, nil
|
|
|
|
case tea.WindowSizeMsg:
|
|
m.progress.Width = msg.Width - 4
|
|
return m, nil
|
|
|
|
case restoreProgressMsg:
|
|
progressVal := float64(msg) / float64(m.totalBytes)
|
|
if progressVal > 1.0 {
|
|
progressVal = 1.0
|
|
}
|
|
|
|
cmd := m.progress.SetPercent(progressVal)
|
|
|
|
return m, tea.Batch(cmd, m.restoreProgressListener)
|
|
|
|
case restoreDoneMsg:
|
|
m.done = true
|
|
cmd := m.progress.SetPercent(1.0)
|
|
return m, tea.Batch(cmd, tea.Quit)
|
|
|
|
case restoreErrorMsg:
|
|
m.err = msg.err
|
|
return m, tea.Quit
|
|
}
|
|
|
|
return m, nil
|
|
}
|
|
|
|
func (m *restoreModel) View() string {
|
|
if m.err != nil {
|
|
return fmt.Sprintf("\n❌ Error: %v\n\n", m.err)
|
|
}
|
|
|
|
if m.done {
|
|
duration := time.Since(m.startTime)
|
|
restoreType := "File/Disk"
|
|
if m.isDirectory {
|
|
restoreType = "Directory"
|
|
}
|
|
return fmt.Sprintf("\n✅ %s restoration complete!\n\n"+
|
|
"Archive: %s → %s\n"+
|
|
"Time: %v\n\n",
|
|
restoreType, m.source, m.target,
|
|
duration.Round(time.Millisecond))
|
|
}
|
|
|
|
title := fmt.Sprintf("Restoring %s → %s", m.source, m.target)
|
|
progressVal := float64(m.processed) / float64(m.totalBytes)
|
|
progressView := m.progress.ViewAs(progressVal)
|
|
stats := fmt.Sprintf("%s / %s (%.1f%%)", formatBytes(m.processed), formatBytes(m.totalBytes), progressVal*100)
|
|
|
|
return lipgloss.JoinVertical(lipgloss.Left,
|
|
title,
|
|
"",
|
|
progressView,
|
|
"",
|
|
stats,
|
|
"\nPress Ctrl+C to cancel",
|
|
)
|
|
}
|
|
|
|
func isBlockDevice(path string) bool {
|
|
if runtime.GOOS != "linux" {
|
|
return false
|
|
}
|
|
|
|
fileInfo, err := os.Stat(path)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
|
|
stat, ok := fileInfo.Sys().(*syscall.Stat_t)
|
|
if !ok {
|
|
return false
|
|
}
|
|
|
|
return (stat.Mode & syscall.S_IFMT) == syscall.S_IFBLK
|
|
}
|
|
|
|
func init() {
|
|
rootCmd.AddCommand(restoreCmd)
|
|
restoreCmd.Flags().StringVarP(&restoreSource, "source", "s", "", "Source archive (required)")
|
|
restoreCmd.Flags().StringVarP(&restoreTarget, "target", "t", "", "Target file, disk, or directory. Defaults to current directory for archives.")
|
|
restoreCmd.MarkFlagRequired("source")
|
|
}
|