HellreigN/backend/internal/handlers/jobs.go

package handlers

import (
	"errors"
	"fmt"
	"net/http"
	"os/exec"
	"strconv"

	"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/backend/internal/grpcsrv/commander"
	"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/backend/internal/models"
	"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/backend/internal/service"
	"github.com/gin-gonic/gin"
)

type JobsHandlers struct {
	tracker  *commander.ConnTracker
	svc      *service.ScriptService
	whereami string
}

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"`
	Stdout  string   `json:"stdout"`
	Stderr  string   `json:"stderr"`
	Status  int32    `json:"status"`
}

// 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
// @Router /jobs [post]
func (h *JobsHandlers) AddJob(c *gin.Context) {
	var in AddJobIn
	if err := c.Bind(&in); err != nil {
		c.Error(err)
		return
	}

	agent, ok := h.tracker.GetAgent(in.AgentID)
	if !ok {
		c.Status(http.StatusNotFound)
		c.Error(fmt.Errorf("agent not found"))
		return
	}

	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"`
}