corrosion-admin-panel/companion-agent/internal/update/updater.go

package update

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"time"
)

// Updater handles self-update operations for the companion agent
type Updater struct {
	currentVersion string
}

// NewUpdater creates a new updater instance
func NewUpdater(currentVersion string) *Updater {
	return &Updater{
		currentVersion: currentVersion,
	}
}

// PerformUpdate downloads a new version and replaces the current binary
func (u *Updater) PerformUpdate(downloadURL, newVersion string) error {
	log.Printf("Performing self-update from %s to %s", u.currentVersion, newVersion)
	log.Printf("Download URL: %s", downloadURL)

	// Get current executable path
	exePath, err := os.Executable()
	if err != nil {
		return fmt.Errorf("failed to get executable path: %w", err)
	}

	log.Printf("Current executable: %s", exePath)

	// Create temporary file for download
	tmpFile := exePath + ".new"

	// Download new binary
	if err := u.downloadBinary(downloadURL, tmpFile); err != nil {
		return fmt.Errorf("failed to download update: %w", err)
	}

	// Make new binary executable (Unix only)
	if runtime.GOOS != "windows" {
		if err := os.Chmod(tmpFile, 0755); err != nil {
			os.Remove(tmpFile)
			return fmt.Errorf("failed to make binary executable: %w", err)
		}
	}

	// Create backup of current binary
	backupFile := exePath + ".backup"
	if err := u.createBackup(exePath, backupFile); err != nil {
		os.Remove(tmpFile)
		return fmt.Errorf("failed to create backup: %w", err)
	}

	log.Println("Backup created successfully")

	// Replace current binary with new one
	if err := u.replaceBinary(tmpFile, exePath); err != nil {
		// Attempt to restore from backup
		log.Printf("Update failed, attempting to restore from backup: %v", err)
		if restoreErr := u.replaceBinary(backupFile, exePath); restoreErr != nil {
			return fmt.Errorf("update failed and backup restoration failed: %w (original error: %v)", restoreErr, err)
		}
		return fmt.Errorf("update failed, restored from backup: %w", err)
	}

	log.Printf("Successfully updated to version %s", newVersion)
	log.Println("NOTE: Restart the agent to use the new version")

	// Clean up backup file after successful update
	os.Remove(backupFile)

	return nil
}

// downloadBinary downloads a binary from the given URL to the destination path
func (u *Updater) downloadBinary(url, destPath string) error {
	log.Printf("Downloading binary from %s", url)

	// Create HTTP client with timeout
	client := &http.Client{
		Timeout: 10 * time.Minute, // Binary download may take longer
	}

	// Download file
	resp, err := client.Get(url)
	if err != nil {
		return fmt.Errorf("failed to download: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("download failed with status: %d", resp.StatusCode)
	}

	// Create destination file
	file, err := os.Create(destPath)
	if err != nil {
		return fmt.Errorf("failed to create file: %w", err)
	}
	defer file.Close()

	// Copy content
	written, err := io.Copy(file, resp.Body)
	if err != nil {
		return fmt.Errorf("failed to write file: %w", err)
	}

	log.Printf("Downloaded %d bytes", written)
	return nil
}

// createBackup creates a backup copy of a file
func (u *Updater) createBackup(src, dest string) error {
	srcFile, err := os.Open(src)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	destFile, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer destFile.Close()

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

// replaceBinary replaces the destination binary with the source binary
func (u *Updater) replaceBinary(src, dest string) error {
	// On Windows, we can't replace a running executable directly
	// We need to rename the old one and move the new one in place
	if runtime.GOOS == "windows" {
		oldExe := dest + ".old"

		// Remove any existing .old file
		os.Remove(oldExe)

		// Rename current executable
		if err := os.Rename(dest, oldExe); err != nil {
			return fmt.Errorf("failed to rename current executable: %w", err)
		}

		// Move new executable into place
		if err := os.Rename(src, dest); err != nil {
			// Try to restore
			os.Rename(oldExe, dest)
			return fmt.Errorf("failed to move new executable: %w", err)
		}

		// Schedule old executable for deletion on next boot
		// (Windows doesn't allow deleting running executables)
		return nil
	}

	// On Unix, we can replace the file directly
	// The running process will continue using the old inode
	return os.Rename(src, dest)
}

// GetCurrentVersion returns the current version
func (u *Updater) GetCurrentVersion() string {
	return u.currentVersion
}

// VerifyUpdate verifies that an update was successful by checking the version
func (u *Updater) VerifyUpdate(expectedVersion string) error {
	exePath, err := os.Executable()
	if err != nil {
		return fmt.Errorf("failed to get executable path: %w", err)
	}

	info, err := os.Stat(exePath)
	if err != nil {
		return fmt.Errorf("failed to stat executable: %w", err)
	}

	log.Printf("Executable info: size=%d, mod_time=%s", info.Size(), info.ModTime())

	// In a real implementation, you might embed version info in the binary
	// or verify a checksum. For now, we just verify the file exists and was modified recently
	if time.Since(info.ModTime()) > 5*time.Minute {
		return fmt.Errorf("executable was not recently modified")
	}

	return nil
}

// CleanupOldVersions removes old backup files
func (u *Updater) CleanupOldVersions() error {
	exePath, err := os.Executable()
	if err != nil {
		return err
	}

	dir := filepath.Dir(exePath)
	baseName := filepath.Base(exePath)

	// Remove backup files
	patterns := []string{
		baseName + ".backup",
		baseName + ".old",
		baseName + ".new",
	}

	for _, pattern := range patterns {
		path := filepath.Join(dir, pattern)
		if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
			log.Printf("Failed to remove old version %s: %v", path, err)
		}
	}

	return nil
}