FemaInstaller/internal/updater/updater.go

package updater

import (
	"fmt"
	"os"
	"path/filepath"
	"time"

	"gitea.unprism.ru/KRBL/FemaInstaller/internal/ssh"
	"gitea.unprism.ru/KRBL/FemaInstaller/pkg/config"
	"gitea.unprism.ru/KRBL/FemaInstaller/pkg/fileutils"
	gossh "golang.org/x/crypto/ssh"
)

// ProgressInfo contains information about the update progress
type ProgressInfo struct {
	Stage                  string
	Percentage             float64
	EstimatedTimeRemaining time.Duration
}

// ProgressCallback is a function that reports update progress
type ProgressCallback func(progress ProgressInfo)

// UpdateMethod defines the method used for updating
type UpdateMethod int

const (
	// UpdateMethodEmbedded uses the binary embedded in the program
	UpdateMethodEmbedded UpdateMethod = iota
	// UpdateMethodDirectDownload uses wget to download directly on the device
	UpdateMethodDirectDownload
)

// Updater handles the software update process
type Updater struct {
	Config       *config.UpdaterConfig
	BinaryData   []byte
	UpdateMethod UpdateMethod
}

// NewUpdater creates a new updater with the provided configuration and binary data
func NewUpdater(config *config.UpdaterConfig, binaryData []byte) *Updater {
	return &Updater{
		Config:       config,
		BinaryData:   binaryData,
		UpdateMethod: UpdateMethodEmbedded, // Default to embedded method
	}
}

// SetUpdateMethod sets the update method to use
func (u *Updater) SetUpdateMethod(method UpdateMethod) {
	u.UpdateMethod = method
}

// Update performs the software update process
func (u *Updater) Update(progressCallback ProgressCallback) error {
	// If no callback is provided, use a no-op callback
	if progressCallback == nil {
		progressCallback = func(progress ProgressInfo) {}
	}

	// Report initial progress
	progressCallback(ProgressInfo{
		Stage:                  "Подготовка",
		Percentage:             0,
		EstimatedTimeRemaining: 0,
	})

	// Create SSH client configuration
	clientConfig := ssh.NewClientConfig(
		u.Config.IP,
		u.Config.Port,
		u.Config.Login,
		u.Config.Password,
	)

	// Connect to SSH server
	progressCallback(ProgressInfo{
		Stage:                  "Подключение к серверу",
		Percentage:             10,
		EstimatedTimeRemaining: 0,
	})

	sshClient, err := ssh.CreateSSHClient(clientConfig)
	if err != nil {
		return fmt.Errorf("ошибка подключения SSH: %w", err)
	}
	defer sshClient.Close()

	// Choose update method
	switch u.UpdateMethod {
	case UpdateMethodEmbedded:
		return u.updateWithEmbeddedBinary(sshClient, clientConfig, progressCallback)
	case UpdateMethodDirectDownload:
		return u.updateWithDirectDownload(sshClient, progressCallback)
	default:
		return fmt.Errorf("неизвестный метод обновления")
	}
}

// updateWithEmbeddedBinary updates using the binary embedded in the program
func (u *Updater) updateWithEmbeddedBinary(sshClient *gossh.Client, clientConfig *ssh.ClientConfig, progressCallback ProgressCallback) error {
	// Save binary to temporary file
	tempFile := "update_binary"
	if err := os.WriteFile(tempFile, u.BinaryData, 0644); err != nil {
		return fmt.Errorf("не удалось сохранить временный файл: %w", err)
	}
	defer os.Remove(tempFile) // Clean up temporary file

	progressCallback(ProgressInfo{
		Stage:                  "Создание SFTP-соединения",
		Percentage:             20,
		EstimatedTimeRemaining: 0,
	})

	// Create SFTP client
	sftpClient, err := ssh.CreateSFTPClient(clientConfig)
	if err != nil {
		return fmt.Errorf("ошибка подключения SFTP: %w", err)
	}
	defer sftpClient.Close()

	progressCallback(ProgressInfo{
		Stage:                  "Загрузка файла",
		Percentage:             30,
		EstimatedTimeRemaining: 0,
	})

	// Upload binary with progress reporting
	remotePath := "/root/fema/build.new"

	// Create a file upload progress callback
	uploadProgressCallback := func(percentage float64, estimatedTimeRemaining time.Duration) {
		// Map the upload progress (0-100%) to the overall progress (30-80%)
		overallPercentage := 30 + (percentage * 0.5)

		progressCallback(ProgressInfo{
			Stage:                  "Загрузка файла",
			Percentage:             overallPercentage,
			EstimatedTimeRemaining: estimatedTimeRemaining,
		})
	}

	if err = fileutils.UploadFileWithProgress(sftpClient, tempFile, remotePath, uploadProgressCallback); err != nil {
		return fmt.Errorf("ошибка загрузки файла: %w", err)
	}

	return u.finalizeUpdate(sshClient, progressCallback)
}

// updateWithDirectDownload updates by downloading directly on the device using wget
func (u *Updater) updateWithDirectDownload(sshClient *gossh.Client, progressCallback ProgressCallback) error {
	if u.Config.DownloadURL == "" {
		return fmt.Errorf("URL загрузки не указан в конфигурации")
	}

	progressCallback(ProgressInfo{
		Stage:                  "Загрузка файла на устройство",
		Percentage:             30,
		EstimatedTimeRemaining: 0,
	})

	// Download file directly on the device using wget
	downloadCmd := fmt.Sprintf("wget -O /root/fema/build.new %s", u.Config.DownloadURL)

	if err := ssh.ExecuteCommand(sshClient, downloadCmd); err != nil {
		return fmt.Errorf("ошибка загрузки файла на устройство: %w", err)
	}

	progressCallback(ProgressInfo{
		Stage:                  "Загрузка файла на устройство",
		Percentage:             80,
		EstimatedTimeRemaining: 0,
	})

	return u.finalizeUpdate(sshClient, progressCallback)
}

// finalizeUpdate applies the update by moving files and restarting the service
func (u *Updater) finalizeUpdate(sshClient *gossh.Client, progressCallback ProgressCallback) error {
	progressCallback(ProgressInfo{
		Stage:                  "Применение обновления",
		Percentage:             80,
		EstimatedTimeRemaining: 0,
	})

	// Execute update commands
	commands := []string{
		"mv -f /root/fema/build.new /root/fema/build",
		"chmod +x /root/fema/build",
		"systemctl restart fema.service",
	}

	for i, cmd := range commands {
		// Calculate progress for each command (80-100%)
		cmdProgress := 80 + float64(i+1)*20/float64(len(commands))

		progressCallback(ProgressInfo{
			Stage:                  "Применение обновления",
			Percentage:             cmdProgress,
			EstimatedTimeRemaining: 0,
		})

		if err := ssh.ExecuteCommand(sshClient, cmd); err != nil {
			return fmt.Errorf("ошибка выполнения команды '%s': %w", cmd, err)
		}
	}

	// Report completion
	progressCallback(ProgressInfo{
		Stage:                  "Завершено",
		Percentage:             100,
		EstimatedTimeRemaining: 0,
	})

	return nil
}

// GetConfigFilePath returns the path to the configuration file
func GetConfigFilePath() string {
	// Get executable directory
	execDir, err := filepath.Abs(filepath.Dir(os.Args[0]))
	if err != nil {
		execDir = "."
	}

	return filepath.Join(execDir, "updater_config.json")
}