backup/cmd/create.go

package cmd

import (
	"archive/tar"
	"compress/gzip"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"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
	excludePatterns []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, excludePatterns)
		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
	httpTargetInfo  *ServerTargetInfo
	isP2PTarget     bool // Флаг для старого P2P-режима
	excludePatterns []string
}

type progressMsg int64
type doneMsg struct{}
type errorMsg struct{ err error }

func getPathSize(path string, excludePatterns []string) (int64, error) {
	fileInfo, err := os.Stat(path)
	if err != nil {
		return 0, err
	}

	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(current string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}

			relPath, err := filepath.Rel(path, current)
			if err != nil {
				return err
			}

			if relPath == "." {
				return nil
			}

			if shouldExcludePath(relPath, excludePatterns) {
				if info.IsDir() {
					return filepath.SkipDir
				}
				return nil
			}

			if !info.IsDir() {
				total += info.Size()
			}
			return nil
		})
		return total, err
	}

	return fileInfo.Size(), nil
}

func initialModel(src, dst string, excludePatterns []string) (*model, error) {
	totalBytes, err := getPathSize(src, excludePatterns)
	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
	}
	isDirectory := srcInfo.IsDir()

	targetInfo, isHttp := parseServerTarget(dst)
	isP2P := false
	targetPath := dst

	if isHttp {
		// Новый HTTP-режим
		if isDirectory && !strings.HasSuffix(targetInfo.Filename, ".tar.gz") {
			targetInfo.Filename += ".tar.gz"
		} else if !isDirectory && !strings.HasSuffix(targetInfo.Filename, ".gz") {
			targetInfo.Filename += ".gz"
		}
		targetInfo.URL = fmt.Sprintf("http://%s/backup/%s", targetInfo.Address, targetInfo.Filename)
		targetPath = targetInfo.URL
	} else if isNetworkAddress(dst) {
		// Старый P2P-режим
		isP2P = true
	} else {
		// Локальный файл
		if isDirectory {
			if !strings.HasSuffix(targetPath, ".tar.gz") {
				targetPath += ".tar.gz"
			}
		} else {
			if !strings.HasSuffix(targetPath, ".gz") {
				targetPath += ".gz"
			}
		}
	}

	return &model{
		progress:        p,
		source:          src,
		target:          targetPath,
		totalBytes:      totalBytes,
		startTime:       time.Now(),
		progressCh:      make(chan int64, 100),
		isDirectory:     isDirectory,
		httpTargetInfo:  targetInfo,
		isP2PTarget:     isP2P,
		excludePatterns: excludePatterns,
	}, nil
}

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

func (m *model) getTargetWriter() (io.WriteCloser, error) {
	if m.httpTargetInfo != nil {
		// Новый HTTP-режим
		pipeReader, pipeWriter := io.Pipe()
		req, err := http.NewRequest("POST", m.httpTargetInfo.URL, pipeReader)
		if err != nil {
			return nil, fmt.Errorf("failed to create http request: %w", err)
		}
		req.Header.Set("Content-Type", "application/octet-stream")
		req.ContentLength = -1 // Stream upload

		go func() {
			client := &http.Client{} // No timeout for uploads
			resp, err := client.Do(req)
			if err != nil {
				pipeWriter.CloseWithError(fmt.Errorf("http request failed: %w", err))
				return
			}
			defer resp.Body.Close()

			if resp.StatusCode != http.StatusOK {
				bodyBytes, _ := io.ReadAll(resp.Body)
				err := fmt.Errorf("server returned non-200 status: %s\n%s", resp.Status, string(bodyBytes))
				pipeWriter.CloseWithError(err)
			}
		}()

		return pipeWriter, nil

	} else if m.isP2PTarget {
		// Старый P2P-режим
		conn, err := net.DialTimeout("tcp", m.target, 10*time.Second)
		if err != nil {
			return nil, fmt.Errorf("failed to connect to %s: %w", m.target, err)
		}

		// Отправляем бинарный заголовок (размер + тип)
		if err := binary.Write(conn, binary.BigEndian, m.totalBytes); err != nil {
			conn.Close()
			return nil, fmt.Errorf("failed to send backup size: %w", err)
		}
		var typeByte byte = 0
		if m.isDirectory {
			typeByte = 1
		}
		if _, err := conn.Write([]byte{typeByte}); err != nil {
			conn.Close()
			return nil, fmt.Errorf("failed to send backup type: %w", err)
		}
		return conn, nil
	}

	// Локальный файл
	return os.Create(m.target)
}

func (m *model) startBackup() tea.Msg {
	go func() {
		targetWriter, err := m.getTargetWriter()
		if err != nil {
			m.err = err
			m.progressCh <- -1
			return
		}
		defer targetWriter.Close()

		var backupErr error
		if m.isDirectory {
			backupErr = m.backupDirectory(targetWriter)
		} else {
			backupErr = m.backupFileOrDisk(targetWriter)
		}

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

		m.progressCh <- -2
	}()

	return nil
}

func (m *model) backupFileOrDisk(w io.Writer) error {
	srcFile, err := os.Open(m.source)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	progressWriter := &progressTracker{Writer: w, progressCh: m.progressCh}
	gzipWriter := gzip.NewWriter(progressWriter)
	defer gzipWriter.Close()

	_, err = io.Copy(gzipWriter, srcFile)
	return err
}

// Вспомогательная функция для проверки исключений
func shouldExcludePath(path string, patterns []string) bool {
	cleanPath := filepath.Clean(path)

	for _, pattern := range patterns {
		// Проверка полного совпадения пути
		if matched, _ := filepath.Match(pattern, cleanPath); matched {
			return true
		}

		// Проверка, что паттерн является префиксом пути (как каталог)
		if strings.HasPrefix(cleanPath, pattern+string(filepath.Separator)) {
			return true
		}

		// Проверка базового имени файла/директории
		base := filepath.Base(cleanPath)
		if matched, _ := filepath.Match(pattern, base); matched {
			return true
		}
	}
	return false
}

func (m *model) shouldExclude(path string) bool {
	return shouldExcludePath(path, m.excludePatterns)
}

func (m *model) backupDirectory(w io.Writer) error {
	progressWriter := &progressTracker{Writer: w, progressCh: m.progressCh}
	gzipWriter := gzip.NewWriter(progressWriter)
	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
		}

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

		if relPath != "." {
			if m.shouldExclude(relPath) {
				if fi.IsDir() {
					return filepath.SkipDir
				}
				return nil // Пропускаем файл
			}
		}
		header, err := tar.FileInfoHeader(fi, 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()

		_, err = io.Copy(tarWriter, srcFile)
		return err
	})
}

func (m *model) progressListener() tea.Msg {
	n := <-m.progressCh
	if n < 0 {
		if n == -1 {
			return errorMsg{m.err}
		}
		return doneMsg{}
	}

	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
		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)
		backupType := "File/disk"
		if m.isDirectory {
			backupType = "Directory"
		}
		if m.httpTargetInfo != nil || m.isP2PTarget {
			backupType += " network"
		}

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

	operation := "Backing up"
	if m.isDirectory {
		operation = "Archiving"
	}
	if m.httpTargetInfo != nil || m.isP2PTarget {
		operation = "Sending"
	}

	title := fmt.Sprintf("%s %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 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 or network address (e.g., /path/to/backup.tar.gz, 127.0.0.1:8080/backup.gz, or 127.0.0.1:8080) (required)",
	)
	createCmd.Flags().StringSliceVarP(
		&excludePatterns,
		"exclude",
		"e",
		[]string{},
		"Patterns to exclude (can be specified multiple times)",
	)
	createCmd.MarkFlagRequired("source")
	createCmd.MarkFlagRequired("target")
}