backup/cmd/create.go

package cmd

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"syscall"
	"time"
	"unsafe"

	"github.com/charmbracelet/bubbles/progress"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/spf13/cobra"
	"golang.org/x/sys/unix"
)

var (
	source string
	target string
)

var createCmd = &cobra.Command{
	Use:   "create",
	Short: "Create a new backup",
	Long:  "Create a new backup of a file, directory, or disk",
	Run: func(cmd *cobra.Command, args []string) {
		model, err := initialModel(source, target)
		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 model 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 progressMsg int64
type doneMsg struct{}
type errorMsg struct{ err error }

// getFileSize возвращает размер файла, диска или директории
func getPathSize(path string) (int64, error) {
	fileInfo, err := os.Stat(path)
	if err != nil {
		return 0, err
	}

	// Для Linux: проверка, является ли источник блочным устройством
	if runtime.GOOS == "linux" {
		stat, ok := fileInfo.Sys().(*syscall.Stat_t)
		if ok && (stat.Mode&syscall.S_IFMT) == syscall.S_IFBLK {
			f, err := os.Open(path)
			if err != nil {
				return 0, err
			}
			defer f.Close()

			var size uint64
			_, _, errno := unix.Syscall(
				unix.SYS_IOCTL,
				f.Fd(),
				unix.BLKGETSIZE64,
				uintptr(unsafe.Pointer(&size)),
			)
			if errno != 0 {
				return 0, fmt.Errorf("ioctl BLKGETSIZE64 failed: %v", errno)
			}
			return int64(size), nil
		}
	}

	// Для директорий
	if fileInfo.IsDir() {
		var total int64
		err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.IsDir() {
				total += info.Size()
			}
			return nil
		})
		return total, err
	}

	// Для обычных файлов
	return fileInfo.Size(), nil
}

func initialModel(src, dst string) (*model, error) {
	totalBytes, err := getPathSize(src)
	if err != nil {
		return nil, err
	}

	p := progress.New(
		progress.WithDefaultGradient(),
		progress.WithWidth(40),
	)

	// Проверяем, является ли источник директорией
	srcInfo, err := os.Stat(src)
	if err != nil {
		return nil, err
	}

	// Добавляем расширение .tar.gz для директорий
	targetPath := dst
	if srcInfo.IsDir() && filepath.Ext(targetPath) != ".tar.gz" {
		targetPath += ".tar.gz"
	}

	return &model{
		progress:    p,
		source:      src,
		target:      targetPath,
		totalBytes:  totalBytes,
		startTime:   time.Now(),
		progressCh:  make(chan int64, 100),
		isDirectory: srcInfo.IsDir(),
	}, nil
}

func (m *model) Init() tea.Cmd {
	return tea.Batch(
		m.startBackup,
		m.progressListener,
	)
}

func (m *model) startBackup() tea.Msg {
	go func() {
		if m.isDirectory {
			err := m.backupDirectory()
			if err != nil {
				m.progressCh <- -1
				return
			}
			m.progressCh <- -2
			return
		}

		// Обработка файлов и дисков
		srcFile, err := os.Open(m.source)
		if err != nil {
			m.progressCh <- -1
			return
		}
		defer srcFile.Close()

		dstFile, err := os.Create(m.target + ".gz")
		if err != nil {
			m.progressCh <- -1
			return
		}
		defer dstFile.Close()

		gzipWriter := gzip.NewWriter(dstFile)
		defer gzipWriter.Close()

		buf := make([]byte, 32*1024) // 32KB buffer
		for {
			n, err := srcFile.Read(buf)
			if n > 0 {
				_, writeErr := gzipWriter.Write(buf[:n])
				if writeErr != nil {
					m.progressCh <- -1
					return
				}
				m.progressCh <- int64(n)
			}

			if err == io.EOF {
				break
			}

			if err != nil {
				m.progressCh <- -1
				return
			}
		}

		m.progressCh <- -2 // Сигнал завершения
	}()

	return nil
}

func (m *model) backupDirectory() error {
	dstFile, err := os.Create(m.target)
	if err != nil {
		return err
	}
	defer dstFile.Close()

	gzipWriter := gzip.NewWriter(dstFile)
	defer gzipWriter.Close()

	tarWriter := tar.NewWriter(gzipWriter)
	defer tarWriter.Close()

	return filepath.Walk(m.source, func(file string, fi os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Создаем заголовок для файла/директории
		header, err := tar.FileInfoHeader(fi, file)
		if err != nil {
			return err
		}

		// Устанавливаем относительный путь
		relPath, err := filepath.Rel(m.source, file)
		if err != nil {
			return err
		}
		header.Name = relPath

		// Пишем заголовок в архив
		if err := tarWriter.WriteHeader(header); err != nil {
			return err
		}

		// Если это не обычный файл, пропускаем запись данных
		if !fi.Mode().IsRegular() {
			return nil
		}

		// Открываем файл для чтения
		srcFile, err := os.Open(file)
		if err != nil {
			return err
		}
		defer srcFile.Close()

		// Копируем содержимое файла в архив
		buf := make([]byte, 32*1024)
		for {
			n, err := srcFile.Read(buf)
			if n > 0 {
				if _, err := tarWriter.Write(buf[:n]); err != nil {
					return err
				}
				m.progressCh <- int64(n)
			}

			if err == io.EOF {
				break
			}

			if err != nil {
				return err
			}
		}

		return nil
	})
}

func (m *model) progressListener() tea.Msg {
	select {
	case n := <-m.progressCh:
		switch {
		case n == -1: // Ошибка
			return errorMsg{fmt.Errorf("backup error")}
		case n == -2: // Завершение
			return doneMsg{}
		default: // Прогресс
			m.processed += n
			return progressMsg(m.processed)
		}
	}
}

func (m *model) 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 progressMsg:
		// Обновляем прогрессбар
		progressVal := float64(msg) / float64(m.totalBytes)
		if progressVal > 1.0 {
			progressVal = 1.0
		}

		// Создаем команду для обновления прогрессбара
		cmd := m.progress.SetPercent(progressVal)

		// Продолжаем слушать обновления
		return m, tea.Batch(cmd, m.progressListener)

	case doneMsg:
		m.done = true
		// Устанавливаем 100% и завершаем
		cmd := m.progress.SetPercent(1.0)
		return m, tea.Batch(cmd, tea.Quit)

	case errorMsg:
		m.err = msg.err
		return m, tea.Quit
	}

	return m, nil
}

func (m *model) View() string {
	if m.err != nil {
		return fmt.Sprintf("\n❌ Error: %v\n\n", m.err)
	}

	if m.done {
		duration := time.Since(m.startTime)
		ratio := float64(m.processed) / float64(m.totalBytes) * 100
		backupType := "file/disk"
		if m.isDirectory {
			backupType = "directory"
		}

		return fmt.Sprintf("\n✅ %s backup complete!\n\n"+
			"Source: %s\n"+
			"Backup: %s\n"+
			"Size: %s → %s (%.1f%%)\n"+
			"Time: %v\n\n",
			backupType, m.source, m.target,
			formatBytes(m.totalBytes), formatBytes(m.processed), ratio,
			duration.Round(time.Millisecond))
	}

	backupType := "Backing up"
	if m.isDirectory {
		backupType = "Archiving"
	}
	title := fmt.Sprintf("%s %s → %s", backupType, 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 formatBytes(b int64) string {
	const unit = 1024
	if b < unit {
		return fmt.Sprintf("%d B", b)
	}
	div, exp := int64(unit), 0
	for n := b / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %ciB", float64(b)/float64(div), "KMGTPE"[exp])
}

func init() {
	rootCmd.AddCommand(createCmd)
	createCmd.Flags().StringVarP(&source, "source", "s", "", "Source file, directory or disk (e.g., /dev/sda) (required)")
	createCmd.Flags().StringVarP(&target, "target", "t", "", "Target backup file (required)")
	createCmd.MarkFlagRequired("source")
	createCmd.MarkFlagRequired("target")
}