feat: big ahh commit

- agent+proto+backend: transfer service status - agent: fix returning empty message on nonzero exit status - backend: refactor collector+commander and handlers dependent on them: implement agent accounting via grpc stats handler
2026-04-04 22:55:39 +03:00
parent b516a54c17
commit 7be99f8e91
12 changed files with 541 additions and 190 deletions
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"os/exec"
+	"strconv"

 	"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/backend/internal/grpcsrv/commander"
 	"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/backend/internal/models"
@@ -13,21 +14,32 @@ import (
 )

 type JobsHandlers struct {
-	cmder *commander.Commander
-	svc   *service.ScriptService
+	tracker  *commander.ConnTracker
+	svc      *service.ScriptService
+	whereami string
 }

-func NewJobsHandlers(cmder *commander.Commander, svc *service.ScriptService) JobsHandlers {
-	return JobsHandlers{cmder: cmder, svc: svc}
+func NewJobsHandlers(tracker *commander.ConnTracker, svc *service.ScriptService, whereami string) JobsHandlers {
+	return JobsHandlers{tracker: tracker, svc: svc, whereami: whereami}
 }

+// AddJobIn is the request body for creating a job.
 type AddJobIn struct {
 	Command       string  `json:"command"        binding:"required"`
 	InterpreterID int64   `json:"interpreter_id"`
 	Stdin         *string `json:"stdin"`
 	AgentID       string  `json:"agent_id"       binding:"required"`
 }
+
+// AddJobOut is the response body for a submitted job.
 type AddJobOut struct {
+	ID      int64    `json:"id"`
+	Command []string `json:"command"`
+	WaitURL string   `json:"wait_url"`
+}
+
+// JobResult is the response body for a completed job.
+type JobResult struct {
 	ID      int64    `json:"id"`
 	Command []string `json:"command"`
 	Stdin   *string  `json:"stdin"`
@@ -36,65 +48,152 @@ type AddJobOut struct {
 	Status  int32    `json:"status"`
 }

-// AddJob creates and executes a job on a target agent.
-// @Summary Create and run a job on an agent
-// @Description Sends a command to the specified agent, waits for execution, and returns the result
+// WaitJobIn is the request body for waiting on a job.
+type WaitJobIn struct {
+	AgentID string `json:"agent_id" binding:"required"`
+}
+
+// AddJob submits a job to an agent and returns a wait_url for the result.
+// @Summary Submit a job to an agent
+// @Description Sends a command to the specified agent and returns a URL to wait for the result
 // @Tags jobs
 // @Accept json
 // @Produce json
 // @Param body body AddJobIn true "Job request"
 // @Success 201 {object} AddJobOut
-// @Security Bearer
 // @Router /jobs [post]
-func (self *JobsHandlers) AddJob(c *gin.Context) {
-	err := func() error {
-		var in AddJobIn
-		if err := c.Bind(&in); err != nil {
-			return err
-		}
-		agent, ok := self.cmder.GetAgent(in.AgentID)
-		if !ok {
-			c.Status(http.StatusNotFound)
-			return fmt.Errorf("agent not found")
-		}
+func (h *JobsHandlers) AddJob(c *gin.Context) {
+	var in AddJobIn
+	if err := c.Bind(&in); err != nil {
+		c.Error(err)
+		return
+	}

-		var command []string
-		if in.InterpreterID == 0 {
-			command = []string{"sh", "-c", in.Command}
-		} else {
-			var err error
-			command, err = self.svc.ResolveCommand(
-				c.Request.Context(),
-				in.InterpreterID,
-				in.Command,
-			)
-			if err != nil {
-				return err
-			}
-		}
+	agent, ok := h.tracker.GetAgent(in.AgentID)
+	if !ok {
+		c.Status(http.StatusNotFound)
+		c.Error(fmt.Errorf("agent not found"))
+		return
+	}

-		jid, err := agent.AddJob(models.JobForInsert{
-			Command: command,
-			Stdin:   in.Stdin,
-		})
-		if err != nil {
-			return err
-		}
-		job, err := agent.WaitJob(jid)
-		if err != nil && !errors.As(err, &exec.ExitError{}) {
-			return err
-		}
-		c.JSON(http.StatusCreated, AddJobOut{
-			ID:      job.ID,
-			Command: job.Command,
-			Stdin:   job.Stdin,
-			Stdout:  job.Stdout,
-			Stderr:  job.Stderr,
-			Status:  job.Status,
-		})
-		return nil
-	}()
+	command, err := resolveCommand(c, h.svc, in.InterpreterID, in.Command)
 	if err != nil {
 		c.Error(err)
+		return
 	}
+
+	jid, err := agent.AddJob(models.JobForInsert{
+		Command: command,
+		Stdin:   in.Stdin,
+	})
+	if err != nil {
+		c.Error(err)
+		return
+	}
+
+	waitURL := fmt.Sprintf("%s/api/v1/jobs/%d/wait", h.whereami, jid)
+
+	c.JSON(http.StatusCreated, AddJobOut{
+		ID:      jid,
+		Command: command,
+		WaitURL: waitURL,
+	})
+}
+
+// WaitJob waits for a submitted job to complete (long-poll).
+// If the job is already done, returns immediately.
+// @Summary Wait for job result
+// @Description Long-polls for a job result. Returns immediately if the job is already finished.
+// @Tags jobs
+// @Accept json
+// @Produce json
+// @Param id path int true "Job ID"
+// @Param body body WaitJobIn true "Agent reference"
+// @Success 200 {object} JobResult
+// @Failure 400 {object} map[string]string
+// @Failure 404 {object} map[string]string
+// @Router /jobs/{id}/wait [post]
+func (h *JobsHandlers) WaitJob(c *gin.Context) {
+	jid, err := strconv.ParseInt(c.Param("id"), 10, 64)
+	if err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid job id"})
+		return
+	}
+
+	var in WaitJobIn
+	if err := c.Bind(&in); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
+		return
+	}
+
+	agent, ok := h.tracker.GetAgent(in.AgentID)
+	if !ok {
+		c.Status(http.StatusNotFound)
+		c.Error(fmt.Errorf("agent not found"))
+		return
+	}
+
+	job, err := agent.WaitJob(jid)
+	if err != nil {
+		c.Error(err)
+		return
+	}
+
+	c.JSON(http.StatusOK, JobResult{
+		ID:      job.ID,
+		Command: job.Command,
+		Stdin:   job.Stdin,
+		Stdout:  job.Stdout,
+		Stderr:  job.Stderr,
+		Status:  job.Status,
+	})
+}
+
+func resolveCommand(c *gin.Context, svc *service.ScriptService, interpID int64, cmd string) ([]string, error) {
+	if interpID == 0 {
+		return []string{"sh", "-c", cmd}, nil
+	}
+
+	command, err := svc.ResolveCommand(c.Request.Context(), interpID, cmd)
+	if err != nil {
+		return nil, err
+	}
+
+	return command, nil
+}
+
+// @Summary Check command path
+// @Description Validates that a command binary exists on the system
+// @Tags jobs
+// @Accept json
+// @Param body body CheckCmdIn true "Command to check"
+// @Success 200 {object} CheckCmdOut
+// @Failure 404 {object} map[string]string
+// @Router /jobs/check_cmd [post]
+func (h *JobsHandlers) CheckCmd(c *gin.Context) {
+	var in struct {
+		Command string `json:"command" binding:"required"`
+	}
+	if err := c.Bind(&in); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request"})
+		return
+	}
+
+	if _, err := exec.LookPath(in.Command); err != nil {
+		if errors.Is(err, exec.ErrNotFound) {
+			c.JSON(http.StatusNotFound, gin.H{"error": "command not found"})
+			return
+		}
+		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
+		return
+	}
+
+	c.JSON(http.StatusOK, CheckCmdOut{Exists: true})
+}
+
+type CheckCmdIn struct {
+	Command string `json:"command" binding:"required" example:"bash"`
+}
+type CheckCmdOut struct {
+	Exists bool `json:"exists"`
 }