Add secure automatic agent updates

2026-08-06 21:01:54 +03:00
parent c618eb018c
commit 12aba383f1
6 changed files with 255 additions and 10 deletions
--- a/agent.go
+++ b/agent.go
@@ -99,11 +99,51 @@ type agentInventory struct {
 	CollectedAt      time.Time
 }

+type AgentUpdateManifest struct {
+	Version string `json:"version"`
+	SHA256  string `json:"sha256"`
+	URL     string `json:"url"`
+}
+
+type agentReportResponse struct {
+	Accepted bool                 `json:"accepted"`
+	Update   *AgentUpdateManifest `json:"update,omitempty"`
+}
+
 var guestInventory struct {
 	sync.Mutex
 	value agentInventory
 }

+var agentBinaryInfo struct {
+	sync.Once
+	manifest AgentUpdateManifest
+	path     string
+	err      error
+}
+
+func currentAgentUpdate() (AgentUpdateManifest, string, error) {
+	agentBinaryInfo.Do(func() {
+		agentBinaryInfo.path, agentBinaryInfo.err = os.Executable()
+		if agentBinaryInfo.err != nil {
+			return
+		}
+		file, err := os.Open(agentBinaryInfo.path)
+		if err != nil {
+			agentBinaryInfo.err = err
+			return
+		}
+		defer file.Close()
+		hash := sha256.New()
+		if _, err = io.Copy(hash, file); err != nil {
+			agentBinaryInfo.err = err
+			return
+		}
+		agentBinaryInfo.manifest = AgentUpdateManifest{Version: version, SHA256: hex.EncodeToString(hash.Sum(nil)), URL: "/api/agent/binary"}
+	})
+	return agentBinaryInfo.manifest, agentBinaryInfo.path, agentBinaryInfo.err
+}
+
 type ManagedAgent struct {
 	ID            string      `json:"id"`
 	Name          string      `json:"name"`
@@ -180,8 +220,7 @@ func (s *Store) EnrollAgent(req agentEnrollmentRequest) (agentConfig, error) {
 }

 func (s *Store) SaveAgentReport(id, secret, remote string, report AgentReport) error {
-	var expected string
-	if s.db.QueryRow(`SELECT secret_hash FROM agents WHERE id=?`, id).Scan(&expected) != nil || !hmac.Equal([]byte(expected), []byte(tokenHash(secret))) {
+	if !s.AuthenticateAgent(id, secret) {
 		return errors.New("агент не авторизован")
 	}
 	data, _ := json.Marshal(report)
@@ -189,6 +228,11 @@ func (s *Store) SaveAgentReport(id, secret, remote string, report AgentReport) e
 	return err
 }

+func (s *Store) AuthenticateAgent(id, secret string) bool {
+	var expected string
+	return s.db.QueryRow(`SELECT secret_hash FROM agents WHERE id=?`, id).Scan(&expected) == nil && hmac.Equal([]byte(expected), []byte(tokenHash(secret)))
+}
+
 func (s *Store) Agents() ([]ManagedAgent, error) {
 	rows, err := s.db.Query(`SELECT id,name,vmid,hostname,version,enrolled_at,last_seen,report_json,remote_address FROM agents ORDER BY name`)
 	if err != nil {
@@ -252,10 +296,33 @@ func runGuestAgent(opts agentOptions) error {
 		req.Header.Set("Authorization", "Bearer "+config.ID+"."+config.Secret)
 		response, requestErr := client.Do(req)
 		if requestErr == nil {
-			io.Copy(io.Discard, response.Body)
+			responseData, _ := io.ReadAll(io.LimitReader(response.Body, 64<<10))
 			response.Body.Close()
+			if response.StatusCode == http.StatusUnauthorized && opts.EnrollmentToken != "" {
+				fresh, enrollmentErr := enrollRemote(config, opts.EnrollmentToken)
+				if enrollmentErr == nil {
+					config = fresh
+					if saveErr := saveAgentConfig(opts.ConfigPath, config); saveErr != nil {
+						return saveErr
+					}
+					continue
+				}
+				requestErr = fmt.Errorf("повторная регистрация: %v", enrollmentErr)
+			}
 			if response.StatusCode >= 300 {
-				requestErr = fmt.Errorf("сервер вернул %s", response.Status)
+				if requestErr == nil {
+					requestErr = fmt.Errorf("сервер вернул %s", response.Status)
+				}
+			}
+			if response.StatusCode == http.StatusOK {
+				var result agentReportResponse
+				if json.Unmarshal(responseData, &result) == nil && result.Update != nil && newerAgentVersion(result.Update.Version, version) {
+					if updateErr := applyAgentUpdate(client, config, *result.Update); updateErr != nil {
+						requestErr = fmt.Errorf("автообновление: %v", updateErr)
+					} else {
+						return nil
+					}
+				}
 			}
 		}
 		if requestErr != nil {
@@ -265,6 +332,96 @@ func runGuestAgent(opts agentOptions) error {
 	}
 }

+func newerAgentVersion(candidate, current string) bool {
+	parse := func(value string) [3]int {
+		value = strings.TrimPrefix(strings.TrimSpace(value), "v")
+		value = strings.SplitN(value, "-", 2)[0]
+		parts := strings.Split(value, ".")
+		var result [3]int
+		for i := 0; i < len(parts) && i < 3; i++ {
+			result[i], _ = strconv.Atoi(parts[i])
+		}
+		return result
+	}
+	if candidate == "" || current == "" || current == "dev" {
+		return false
+	}
+	next, installed := parse(candidate), parse(current)
+	for i := 0; i < 3; i++ {
+		if next[i] > installed[i] {
+			return true
+		}
+		if next[i] < installed[i] {
+			return false
+		}
+	}
+	return false
+}
+
+func applyAgentUpdate(client *http.Client, config agentConfig, update AgentUpdateManifest) error {
+	if !strings.HasPrefix(update.URL, "/api/agent/") || len(update.SHA256) != 64 {
+		return errors.New("сервер вернул некорректный манифест")
+	}
+	request, err := http.NewRequest(http.MethodGet, config.Server+update.URL, nil)
+	if err != nil {
+		return err
+	}
+	request.Header.Set("Authorization", "Bearer "+config.ID+"."+config.Secret)
+	updateClient := *client
+	updateClient.Timeout = 2 * time.Minute
+	response, err := updateClient.Do(request)
+	if err != nil {
+		return err
+	}
+	defer response.Body.Close()
+	if response.StatusCode != http.StatusOK {
+		return fmt.Errorf("загрузка вернула %s", response.Status)
+	}
+	executable, err := os.Executable()
+	if err != nil {
+		return err
+	}
+	temporary := executable + ".update"
+	_ = os.Remove(temporary)
+	file, err := os.OpenFile(temporary, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0755)
+	if err != nil {
+		return err
+	}
+	keep := false
+	defer func() {
+		file.Close()
+		if !keep {
+			_ = os.Remove(temporary)
+		}
+	}()
+	hash := sha256.New()
+	written, err := io.Copy(io.MultiWriter(file, hash), io.LimitReader(response.Body, (64<<20)+1))
+	if err != nil {
+		return err
+	}
+	if written > 64<<20 {
+		return errors.New("бинарник превышает допустимый размер")
+	}
+	if err = file.Sync(); err != nil {
+		return err
+	}
+	if err = file.Close(); err != nil {
+		return err
+	}
+	actual := hex.EncodeToString(hash.Sum(nil))
+	if !hmac.Equal([]byte(actual), []byte(strings.ToLower(update.SHA256))) {
+		return errors.New("SHA-256 обновления не совпал")
+	}
+	if err = os.Chmod(temporary, 0755); err != nil {
+		return err
+	}
+	if err = os.Rename(temporary, executable); err != nil {
+		return err
+	}
+	keep = true
+	return nil
+}
+
 func enrollRemote(config agentConfig, token string) (agentConfig, error) {
 	hostname, _ := os.Hostname()
 	machineID, _ := os.ReadFile("/etc/machine-id")