edubox/agent/update.go

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"
)

const updateCheckInterval = 15 * time.Minute

// AgentVersionInfo matches the server's /api/agent/version response.
type AgentVersionInfo struct {
	Version     string `json:"version"`
	DownloadUrls struct {
		Windows    string `json:"windows"`
		WindowsZip string `json:"windowsZip"`
		Linux      string `json:"linux"`
		Mac        string `json:"mac"`
	} `json:"downloadUrls"`
}

// httpBaseURL converts a WebSocket server URL into the corresponding HTTP(S)
// base URL, stripping the /api/websocket path if present.
func httpBaseURL(serverURL string) string {
	u := serverURL
	switch {
	case strings.HasPrefix(u, "wss://"):
		u = "https://" + strings.TrimPrefix(u, "wss://")
	case strings.HasPrefix(u, "ws://"):
		u = "http://" + strings.TrimPrefix(u, "ws://")
	}
	u = strings.TrimSuffix(u, "/api/websocket/")
	u = strings.TrimSuffix(u, "/api/websocket")
	return strings.TrimSuffix(u, "/")
}

// checkForUpdate fetches the latest agent version from the server and compares
// it with the running binary's version.
func checkForUpdate(cfg *AgentConfig) (*AgentVersionInfo, bool, error) {
	if cfg == nil || cfg.Server == "" {
		return nil, false, fmt.Errorf("no server URL configured")
	}
	url := httpBaseURL(cfg.Server) + "/api/agent/version"
	client := httpClientWithProxy(cfg)
	client.Timeout = 30 * time.Second
	resp, err := client.Get(url)
	if err != nil {
		return nil, false, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, false, fmt.Errorf("server returned %d", resp.StatusCode)
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, false, err
	}
	var info AgentVersionInfo
	if err := json.Unmarshal(body, &info); err != nil {
		return nil, false, err
	}
	if info.Version == "" {
		return nil, false, fmt.Errorf("server returned empty version")
	}
	available := info.Version != version
	return &info, available, nil
}

// downloadUpdate downloads the new agent binary to the update directory.
func downloadUpdate(cfg *AgentConfig, dataDir, downloadURL string) (string, error) {
	updateDir := filepath.Join(dataDir, "update")
	if err := os.MkdirAll(updateDir, 0755); err != nil {
		return "", err
	}
	ext := ""
	if runtime.GOOS == "windows" {
		ext = ".exe"
	}
	dest := filepath.Join(updateDir, "studioE5-agent-new"+ext)
	log.Printf("Downloading update from %s to %s", downloadURL, dest)

	client := httpClientWithProxy(cfg)
	resp, err := client.Get(downloadURL)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("download returned %d", resp.StatusCode)
	}
	out, err := os.Create(dest)
	if err != nil {
		return "", err
	}
	defer out.Close()
	if _, err := io.Copy(out, resp.Body); err != nil {
		return "", err
	}
	if err := out.Close(); err != nil {
		return "", err
	}
	if runtime.GOOS != "windows" {
		if err := os.Chmod(dest, 0755); err != nil {
			return "", err
		}
	}
	return dest, nil
}

// formatArgsForShell returns the given arguments as a safely quoted string
// suitable for embedding in shell/PowerShell scripts.
func formatArgsForShell(args []string) string {
	if len(args) == 0 {
		return ""
	}
	quoted := make([]string, len(args))
	for i, a := range args {
		quoted[i] = strconv.Quote(a)
	}
	return strings.Join(quoted, " ")
}

// applyUpdate replaces the running binary with the downloaded one using an
// external helper script, then exits the current process. The new process is
// started with the same arguments as the current one so that tray/console mode
// is preserved.
func applyUpdate(currentPath, newPath, dataDir string) error {
	pid := os.Getpid()
	restartArgs := os.Args[1:]
	switch runtime.GOOS {
	case "windows":
		return applyUpdateWindows(currentPath, newPath, dataDir, pid, restartArgs)
	default:
		return applyUpdateUnix(currentPath, newPath, dataDir, pid, restartArgs)
	}
}

func applyUpdateWindows(currentPath, newPath, dataDir string, pid int, restartArgs []string) error {
	scriptPath := filepath.Join(dataDir, "update", "apply-update.ps1")
	argsList := formatArgsForShell(restartArgs)
	if argsList == "" {
		argsList = ""
	} else {
		argsList = "$startArgs = @(" + argsList + ")"
	}
	script := fmt.Sprintf(`$old = "%s"
$new = "%s"
$targetPid = %d
%s
Wait-Process -Id $targetPid -ErrorAction SilentlyContinue
Start-Sleep -Seconds 2
Move-Item -Path $new -Destination $old -Force
if ($startArgs) {
    Start-Process -FilePath $old -ArgumentList $startArgs -WindowStyle Hidden
} else {
    Start-Process -FilePath $old -WindowStyle Hidden
}
`, currentPath, newPath, pid, argsList)
	if err := os.WriteFile(scriptPath, []byte(script), 0644); err != nil {
		return err
	}
	cmd := exec.Command("powershell.exe", "-ExecutionPolicy", "Bypass", "-WindowStyle", "Hidden", "-File", scriptPath)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	hideWindow(cmd)
	if err := cmd.Start(); err != nil {
		return err
	}
	log.Printf("Update helper started, exiting current process")
	os.Exit(0)
	return nil
}

func applyUpdateUnix(currentPath, newPath, dataDir string, pid int, restartArgs []string) error {
	scriptPath := filepath.Join(dataDir, "update", "apply-update.sh")
	argsList := formatArgsForShell(restartArgs)
	script := fmt.Sprintf(`#!/bin/bash
set -e
old="%s"
new="%s"
pid=%d
while kill -0 "$pid" 2>/dev/null; do sleep 1; done
sleep 2
mv "$new" "$old"
chmod +x "$old"
nohup "$old" %s >/dev/null 2>&1 &
`, currentPath, newPath, pid, argsList)
	if err := os.WriteFile(scriptPath, []byte(script), 0755); err != nil {
		return err
	}
	cmd := exec.Command("/bin/bash", scriptPath)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		return err
	}
	log.Printf("Update helper started, exiting current process")
	os.Exit(0)
	return nil
}

// startAgentUpdate performs the full update flow: download + replace + restart.
func startAgentUpdate(cfg *AgentConfig, dataDir string) error {
	info, available, err := checkForUpdate(cfg)
	if err != nil {
		return fmt.Errorf("update check failed: %w", err)
	}
	if !available {
		return fmt.Errorf("no update available")
	}
	currentPath, err := os.Executable()
	if err != nil {
		return err
	}
	currentPath, err = filepath.Abs(currentPath)
	if err != nil {
		return err
	}
	var downloadURL string
	switch runtime.GOOS {
	case "windows":
		downloadURL = info.DownloadUrls.Windows
	case "darwin":
		downloadURL = info.DownloadUrls.Mac
	default:
		downloadURL = info.DownloadUrls.Linux
	}
	if downloadURL == "" {
		return fmt.Errorf("no download URL for %s", runtime.GOOS)
	}
	newPath, err := downloadUpdate(cfg, dataDir, downloadURL)
	if err != nil {
		return fmt.Errorf("download failed: %w", err)
	}
	log.Printf("Applying update to version %s", info.Version)
	broadcastUI(map[string]interface{}{
		"action":  "update_progress",
		"percent": "90",
		"message": "Redémarrage de l'agent...",
	})
	return applyUpdate(currentPath, newPath, dataDir)
}

// updateCheckerLoop periodically checks for agent updates and notifies the UI.
func updateCheckerLoop(cfg *AgentConfig, dataDir string) {
	for {
		info, available, err := checkForUpdate(cfg)
		if err == nil && available && info != nil {
			log.Printf("Agent update available: %s (current: %s)", info.Version, version)
			setServerAgentVersion(info.Version)
			broadcastUI(map[string]interface{}{
				"action":           "update_available",
				"version":          info.Version,
				"update_available": true,
			})
		}
		time.Sleep(updateCheckInterval)
	}
}