392 lines
8.9 KiB
Go
392 lines
8.9 KiB
Go
package cmd
|
|
|
|
import (
|
|
"archive/tar"
|
|
"compress/gzip"
|
|
"encoding/binary"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"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 a local archive or network stream 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
|
|
isNetworkSource bool
|
|
statusMessage string
|
|
}
|
|
|
|
type restoreProgressMsg int64
|
|
type restoreDoneMsg struct{}
|
|
type restoreErrorMsg struct{ err error }
|
|
|
|
func initialRestoreModel(src, dst string) (*restoreModel, error) {
|
|
isNetwork := isNetworkAddress(src)
|
|
p := progress.New(
|
|
progress.WithDefaultGradient(),
|
|
progress.WithWidth(40),
|
|
)
|
|
|
|
var totalBytes int64 = 1 // Placeholder for network source, to avoid division by zero.
|
|
var isDir bool
|
|
targetPath := dst
|
|
status := "Initializing..."
|
|
|
|
if isNetwork {
|
|
if targetPath == "" {
|
|
return nil, fmt.Errorf("target flag -t is required for a network source")
|
|
}
|
|
status = fmt.Sprintf("Listening on %s...", src)
|
|
} else {
|
|
fileInfo, err := os.Stat(src)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to stat source %s: %w", src, err)
|
|
}
|
|
totalBytes = fileInfo.Size()
|
|
isDir = strings.HasSuffix(src, ".tar.gz")
|
|
|
|
if isDir && targetPath == "" {
|
|
targetPath, err = os.Getwd()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get current directory: %w", err)
|
|
}
|
|
} else if !isDir && targetPath == "" {
|
|
return nil, fmt.Errorf("target flag -t is required for file or disk restoration")
|
|
}
|
|
}
|
|
|
|
return &restoreModel{
|
|
progress: p,
|
|
source: src,
|
|
target: targetPath,
|
|
totalBytes: totalBytes,
|
|
startTime: time.Now(),
|
|
progressCh: make(chan int64, 100),
|
|
isDirectory: isDir,
|
|
isNetworkSource: isNetwork,
|
|
statusMessage: status,
|
|
}, nil
|
|
}
|
|
|
|
func (m *restoreModel) Init() tea.Cmd {
|
|
return tea.Batch(
|
|
m.startRestore,
|
|
m.restoreProgressListener,
|
|
)
|
|
}
|
|
|
|
func (m *restoreModel) getSourceReader() (io.ReadCloser, error) {
|
|
if m.isNetworkSource {
|
|
listener, err := net.Listen("tcp", m.source)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to listen on %s: %w", m.source, err)
|
|
}
|
|
defer listener.Close() // Close listener after accepting one connection
|
|
|
|
m.statusMessage = fmt.Sprintf("Waiting for connection on %s", m.source)
|
|
// This is a blocking call
|
|
conn, err := listener.Accept()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to accept connection: %w", err)
|
|
}
|
|
m.statusMessage = "Client connected. Receiving data..."
|
|
|
|
// Protocol: Read 8 bytes for size, 1 byte for type
|
|
if err := binary.Read(conn, binary.BigEndian, &m.totalBytes); err != nil {
|
|
conn.Close()
|
|
return nil, fmt.Errorf("failed to read backup size: %w", err)
|
|
}
|
|
if m.totalBytes == 0 {
|
|
m.totalBytes = 1 // Avoid division by zero for empty files
|
|
}
|
|
|
|
typeByte := make([]byte, 1)
|
|
if _, err := io.ReadFull(conn, typeByte); err != nil {
|
|
conn.Close()
|
|
return nil, fmt.Errorf("failed to read backup type: %w", err)
|
|
}
|
|
m.isDirectory = (typeByte[0] == 1)
|
|
|
|
return conn, nil
|
|
}
|
|
|
|
// Logic for local file
|
|
m.statusMessage = fmt.Sprintf("Opening archive %s", m.source)
|
|
return os.Open(m.source)
|
|
}
|
|
|
|
func (m *restoreModel) startRestore() tea.Msg {
|
|
go func() {
|
|
sourceReader, err := m.getSourceReader()
|
|
if err != nil {
|
|
m.err = err
|
|
m.progressCh <- -1
|
|
return
|
|
}
|
|
defer sourceReader.Close()
|
|
|
|
progressReader := &restoreProgressReader{
|
|
reader: sourceReader,
|
|
progressCh: m.progressCh,
|
|
}
|
|
|
|
var restoreErr error
|
|
if m.isDirectory {
|
|
restoreErr = m.restoreDirectory(progressReader)
|
|
} else {
|
|
restoreErr = m.restoreFileOrDisk(progressReader)
|
|
}
|
|
|
|
if restoreErr != nil {
|
|
m.err = restoreErr
|
|
m.progressCh <- -1
|
|
return
|
|
}
|
|
|
|
m.progressCh <- -2 // Done signal
|
|
}()
|
|
return nil
|
|
}
|
|
|
|
func (m *restoreModel) restoreFileOrDisk(r io.Reader) error {
|
|
gzipReader, err := gzip.NewReader(r)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer gzipReader.Close()
|
|
|
|
var dstFile *os.File
|
|
if isBlockDevice(m.target) {
|
|
dstFile, err = os.OpenFile(m.target, os.O_WRONLY|os.O_SYNC, 0)
|
|
} else {
|
|
if err := os.MkdirAll(filepath.Dir(m.target), 0755); err != nil {
|
|
return err
|
|
}
|
|
dstFile, err = os.Create(m.target)
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer dstFile.Close()
|
|
|
|
if _, err = io.Copy(dstFile, gzipReader); err != nil {
|
|
return err
|
|
}
|
|
return dstFile.Sync()
|
|
}
|
|
|
|
func (m *restoreModel) restoreDirectory(r io.Reader) error {
|
|
gzipReader, err := gzip.NewReader(r)
|
|
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
|
|
}
|
|
_, copyErr := io.Copy(outFile, tarReader)
|
|
closeErr := outFile.Close()
|
|
if copyErr != nil {
|
|
return copyErr
|
|
}
|
|
if closeErr != nil {
|
|
return closeErr
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type restoreProgressReader struct {
|
|
reader io.Reader
|
|
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: // Ошибка
|
|
if m.err == nil {
|
|
m.err = fmt.Errorf("unknown restoration error")
|
|
}
|
|
return restoreErrorMsg{m.err}
|
|
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"+
|
|
"Source: %s\n"+
|
|
"Target: %s\n"+
|
|
"Time: %v\n\n",
|
|
restoreType, m.source, m.target,
|
|
duration.Round(time.Millisecond))
|
|
}
|
|
|
|
// Show status message before progress starts
|
|
if m.processed == 0 && (m.isNetworkSource || m.statusMessage != "") {
|
|
return m.statusMessage + "\n\nPress Ctrl+C to cancel"
|
|
}
|
|
|
|
title := fmt.Sprintf("Restoring from %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 or network address (e.g., /path/to/archive.tar.gz or 0.0.0.0:8080) (required)")
|
|
restoreCmd.Flags().StringVarP(&restoreTarget, "target", "t", "", "Target file, disk, or directory. Required for network sources.")
|
|
restoreCmd.MarkFlagRequired("source")
|
|
}
|