ProxmoxDash/agent_actions.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"os/exec"
	"regexp"
	"strings"
	"time"
)

type AgentCommand struct {
	ID     int64  `json:"id"`
	Action string `json:"action"`
	Target string `json:"target,omitempty"`
}

type AgentActionResult struct {
	ID      int64  `json:"id"`
	Success bool   `json:"success"`
	Message string `json:"message"`
}

type AgentAction struct {
	ID          int64  `json:"id"`
	AgentID     string `json:"agentId"`
	AgentName   string `json:"agentName"`
	Action      string `json:"action"`
	Target      string `json:"target"`
	RequestedAt int64  `json:"requestedAt"`
	StartedAt   *int64 `json:"startedAt,omitempty"`
	CompletedAt *int64 `json:"completedAt,omitempty"`
	Status      string `json:"status"`
	Result      string `json:"result"`
}

var safeUnitName = regexp.MustCompile(`^[A-Za-z0-9_.@:-]+\.service$`)
var safeContainerName = regexp.MustCompile(`^[A-Za-z0-9_.-]{1,128}$`)

func (s *Store) QueueAgentAction(agentID, action, target string) (AgentAction, error) {
	action, target = strings.TrimSpace(action), strings.TrimSpace(target)
	var raw, name string
	if err := s.db.QueryRow(`SELECT report_json,name FROM agents WHERE id=?`, agentID).Scan(&raw, &name); err != nil {
		return AgentAction{}, fmt.Errorf("агент не найден")
	}
	var report AgentReport
	_ = json.Unmarshal([]byte(raw), &report)
	valid := false
	switch action {
	case "systemd-restart", "systemd-reset-failed":
		if safeUnitName.MatchString(target) {
			for _, service := range report.Services {
				if service.Name == target {
					valid = true
					break
				}
			}
		}
	case "docker-restart":
		if safeContainerName.MatchString(target) {
			for _, container := range report.Containers {
				if container.ID == target || container.Name == target {
					valid = true
					break
				}
			}
		}
	case "updates-refresh":
		valid = target == ""
	}
	if !valid {
		return AgentAction{}, fmt.Errorf("действие или цель не входят в разрешённый список")
	}
	now := time.Now().Unix()
	result, err := s.db.Exec(`INSERT INTO agent_actions(agent_id,action,target,requested_at,status) VALUES(?,?,?,?, 'queued')`, agentID, action, target, now)
	if err != nil {
		return AgentAction{}, err
	}
	id, _ := result.LastInsertId()
	return AgentAction{ID: id, AgentID: agentID, AgentName: name, Action: action, Target: target, RequestedAt: now, Status: "queued"}, nil
}

func (s *Store) DispatchAgentActions(agentID string) ([]AgentCommand, error) {
	tx, err := s.db.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()
	_, _ = tx.Exec(`UPDATE agent_actions SET status='queued',started_at=NULL WHERE agent_id=? AND status='running' AND started_at<?`, agentID, time.Now().Add(-5*time.Minute).Unix())
	rows, err := tx.Query(`SELECT id,action,target FROM agent_actions WHERE agent_id=? AND status='queued' ORDER BY id LIMIT 3`, agentID)
	if err != nil {
		return nil, err
	}
	var commands []AgentCommand
	for rows.Next() {
		var c AgentCommand
		if err = rows.Scan(&c.ID, &c.Action, &c.Target); err != nil {
			rows.Close()
			return nil, err
		}
		commands = append(commands, c)
	}
	rows.Close()
	now := time.Now().Unix()
	for _, c := range commands {
		if _, err = tx.Exec(`UPDATE agent_actions SET status='running',started_at=? WHERE id=?`, now, c.ID); err != nil {
			return nil, err
		}
	}
	return commands, tx.Commit()
}

func (s *Store) CompleteAgentActions(agentID string, results []AgentActionResult) {
	now := time.Now().Unix()
	for _, result := range results {
		status := "failed"
		if result.Success {
			status = "completed"
		}
		message := result.Message
		if len(message) > 4000 {
			message = message[:4000]
		}
		_, _ = s.db.Exec(`UPDATE agent_actions SET status=?,result=?,completed_at=? WHERE id=? AND agent_id=? AND status='running'`, status, message, now, result.ID, agentID)
	}
}

func (s *Store) AgentActions(limit int) ([]AgentAction, error) {
	if limit < 1 || limit > 200 {
		limit = 100
	}
	rows, err := s.db.Query(`SELECT a.id,a.agent_id,g.name,a.action,a.target,a.requested_at,a.started_at,a.completed_at,a.status,a.result FROM agent_actions a JOIN agents g ON g.id=a.agent_id ORDER BY a.id DESC LIMIT ?`, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []AgentAction
	for rows.Next() {
		var v AgentAction
		if err = rows.Scan(&v.ID, &v.AgentID, &v.AgentName, &v.Action, &v.Target, &v.RequestedAt, &v.StartedAt, &v.CompletedAt, &v.Status, &v.Result); err != nil {
			return nil, err
		}
		out = append(out, v)
	}
	return out, rows.Err()
}

func executeAgentCommand(command AgentCommand) AgentActionResult {
	result := AgentActionResult{ID: command.ID}
	var output string
	switch command.Action {
	case "systemd-restart":
		output = runAgentCombinedCommand(30*time.Second, "systemctl", "restart", command.Target)
	case "systemd-reset-failed":
		output = runAgentCombinedCommand(15*time.Second, "systemctl", "reset-failed", command.Target)
	case "docker-restart":
		output = restartAgentContainer(command.Target)
	case "updates-refresh":
		guestInventory.Lock()
		guestInventory.value.CollectedAt = time.Time{}
		guestInventory.Unlock()
		result.Success = true
		result.Message = "Проверка обновлений поставлена в очередь"
		return result
	default:
		result.Message = "Команда не разрешена"
		return result
	}
	result.Success = !strings.HasPrefix(output, "ERROR: ")
	result.Message = strings.TrimSpace(output)
	if result.Message == "" && result.Success {
		result.Message = "Выполнено успешно"
	}
	return result
}

func runAgentCombinedCommand(timeout time.Duration, name string, args ...string) string {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	output, err := exec.CommandContext(ctx, name, args...).CombinedOutput()
	message := strings.TrimSpace(string(output))
	if err != nil {
		return "ERROR: " + err.Error() + ": " + message
	}
	return message
}

func restartAgentContainer(target string) string {
	transport := &http.Transport{DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
		return (&net.Dialer{}).DialContext(ctx, "unix", "/var/run/docker.sock")
	}}
	defer transport.CloseIdleConnections()
	client := &http.Client{Transport: transport, Timeout: 35 * time.Second}
	request, _ := http.NewRequest(http.MethodPost, "http://docker/containers/"+target+"/restart?t=10", bytes.NewReader(nil))
	response, err := client.Do(request)
	if err != nil {
		return "ERROR: " + err.Error()
	}
	defer response.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(response.Body, 4096))
	if response.StatusCode >= 300 {
		return "ERROR: Docker вернул " + response.Status + ": " + strings.TrimSpace(string(body))
	}
	return "Контейнер перезапущен"
}