package cmd import ( "archive/tar" "compress/gzip" "encoding/binary" "fmt" "io" "net" "net/http" "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, P2P stream, or HTTP server.", 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 httpSourceInfo *ServerTargetInfo isP2PSource bool statusMessage string } // restoreSourceReadyMsg is sent when the source is ready to be read. // It carries the reader, the total size, and whether it's a directory backup. type restoreSourceReadyMsg struct { reader io.ReadCloser totalBytes int64 isDir bool } type restoreProgressMsg int64 type restoreDoneMsg struct{} type restoreErrorMsg struct{ err error } func initialRestoreModel(src, dst string) (*restoreModel, error) { p := progress.New( progress.WithDefaultGradient(), progress.WithWidth(40), ) var totalBytes int64 = 1 // Placeholder, will be updated. var isDir, isP2P bool targetPath := dst status := "Initializing..." httpInfo, isHttp := parseServerTarget(src) if isHttp { isDir = strings.HasSuffix(httpInfo.Filename, ".tar.gz") status = fmt.Sprintf("Connecting to %s...", httpInfo.Address) src = httpInfo.URL } else if isNetworkAddress(src) { isP2P = true 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 == "" { // For directory restores, we can default to the current directory. // We do this check here, but also again after getting the P2P type. var err error targetPath, err = os.Getwd() if err != nil { return nil, fmt.Errorf("failed to get current directory: %w", err) } } else if !isDir && targetPath == "" && !isP2P { // Target is mandatory for file/disk restores (except P2P, where we can't know the type yet). return nil, fmt.Errorf("target flag -t is required for this restoration") } return &restoreModel{ progress: p, source: src, target: targetPath, totalBytes: totalBytes, startTime: time.Now(), progressCh: make(chan int64, 100), isDirectory: isDir, httpSourceInfo: httpInfo, isP2PSource: isP2P, statusMessage: status, }, nil } func (m *restoreModel) Init() tea.Cmd { return m.fetchSourceCmd } func (m *restoreModel) fetchSourceCmd() tea.Msg { if m.httpSourceInfo != nil { // New HTTP Mode client := &http.Client{} resp, err := client.Get(m.source) if err != nil { return restoreErrorMsg{fmt.Errorf("failed to connect to server: %w", err)} } if resp.StatusCode != http.StatusOK { bodyBytes, _ := io.ReadAll(resp.Body) resp.Body.Close() return restoreErrorMsg{fmt.Errorf("server returned non-200 status: %s\n%s", resp.Status, string(bodyBytes))} } return restoreSourceReadyMsg{reader: resp.Body, totalBytes: resp.ContentLength, isDir: m.isDirectory} } else if m.isP2PSource { // Old P2P Mode listener, err := net.Listen("tcp", m.source) if err != nil { return restoreErrorMsg{fmt.Errorf("failed to listen on %s: %w", m.source, err)} } // Accept is blocking, so it must be in a command. conn, err := listener.Accept() listener.Close() // Close listener after one connection. if err != nil { return restoreErrorMsg{fmt.Errorf("failed to accept connection: %w", err)} } var size int64 if err := binary.Read(conn, binary.BigEndian, &size); err != nil { conn.Close() return restoreErrorMsg{fmt.Errorf("failed to read backup size: %w", err)} } typeByte := make([]byte, 1) if _, err := io.ReadFull(conn, typeByte); err != nil { conn.Close() return restoreErrorMsg{fmt.Errorf("failed to read backup type: %w", err)} } return restoreSourceReadyMsg{reader: conn, totalBytes: size, isDir: typeByte[0] == 1} } // Local File Mode fileInfo, err := os.Stat(m.source) if err != nil { return restoreErrorMsg{err} } file, err := os.Open(m.source) if err != nil { return restoreErrorMsg{err} } return restoreSourceReadyMsg{reader: file, totalBytes: fileInfo.Size(), isDir: m.isDirectory} } func (m *restoreModel) processRestoreCmd(r io.ReadCloser) tea.Cmd { return func() tea.Msg { defer r.Close() progressReader := &progressTracker{ Reader: r, 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 } else { 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() _, err = io.Copy(dstFile, gzipReader) if err != nil && err != io.EOF { 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 } func (m *restoreModel) restoreProgressListener() tea.Msg { n := <-m.progressCh if n < 0 { if n == -1 { if m.err == nil { m.err = fmt.Errorf("unknown restoration error") } return restoreErrorMsg{m.err} } return restoreDoneMsg{} } 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 restoreSourceReadyMsg: m.totalBytes = msg.totalBytes if m.totalBytes <= 0 { m.totalBytes = 1 } m.isDirectory = msg.isDir // Update directory type, especially for P2P. // Final check for target directory now that we know the type for sure. if m.isDirectory && m.target == "" { wd, err := os.Getwd() if err != nil { return m, func() tea.Msg { return restoreErrorMsg{err} } } m.target = wd } else if !m.isDirectory && m.target == "" { err := fmt.Errorf("target flag -t is required for file/disk restoration") return m, func() tea.Msg { return restoreErrorMsg{err} } } m.statusMessage = "Receiving data..." return m, tea.Batch( m.processRestoreCmd(msg.reader), m.restoreProgressListener, ) 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" } if m.httpSourceInfo != nil || m.isP2PSource { restoreType += " network" } 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)) } if m.processed == 0 { return m.statusMessage + "\n\nPress Ctrl+C to cancel" } operation := "Restoring" if m.httpSourceInfo != nil { operation = "Downloading" } else if m.isP2PSource { operation = "Receiving" } title := fmt.Sprintf("%s from %s → %s", operation, 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, P2P address (e.g., 0.0.0.0:8080), or HTTP URL (e.g., 127.0.0.1:8080/backup.gz) (required)") restoreCmd.Flags().StringVarP(&restoreTarget, "target", "t", "", "Target file, disk, or directory. Required for non-directory restores.") restoreCmd.MarkFlagRequired("source") }