d3m0k1d
275 lines
7.8 KiB
Go
275 lines
7.8 KiB
Go
package handlers
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"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/repository"
|
|
"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/backend/internal/service"
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// ScriptHandlersGroup handles script management routes.
|
|
type ScriptHandlersGroup struct {
|
|
svc *service.ScriptService
|
|
cmder *commander.Commander
|
|
}
|
|
|
|
// NewScriptHandlersGroup creates a new ScriptHandlersGroup.
|
|
func NewScriptHandlersGroup(svc *service.ScriptService, cmder *commander.Commander) *ScriptHandlersGroup {
|
|
return &ScriptHandlersGroup{svc: svc, cmder: cmder}
|
|
}
|
|
|
|
// GetTree returns the script directory tree.
|
|
// @Summary Get script directory tree
|
|
// @Description Returns a hierarchical tree of all scripts organized by their paths
|
|
// @Tags scripts
|
|
// @Produce json
|
|
// @Success 200 {array} repository.ScriptTreeNode
|
|
// @Security Bearer
|
|
// @Router /scripts/tree [get]
|
|
func (sh *ScriptHandlersGroup) GetTree(c *gin.Context) {
|
|
tree, err := sh.svc.BuildTree()
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to build script tree"})
|
|
return
|
|
}
|
|
|
|
if tree == nil {
|
|
tree = []repository.ScriptTreeNode{}
|
|
}
|
|
|
|
c.JSON(http.StatusOK, tree)
|
|
}
|
|
|
|
// CreateScript creates a new script.
|
|
// @Summary Create script
|
|
// @Description Creates a new script with path, content, and interpreter binding
|
|
// @Tags scripts
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Param body body repository.ScriptCreate true "Script data"
|
|
// @Success 201 {object} repository.Script
|
|
// @Security Bearer
|
|
// @Router /scripts [post]
|
|
func (sh *ScriptHandlersGroup) CreateScript(c *gin.Context) {
|
|
var req repository.ScriptCreate
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
|
|
return
|
|
}
|
|
|
|
script, err := sh.svc.Repo.CreateScript(req)
|
|
if err != nil {
|
|
if isUniqueConstraint(err) {
|
|
c.JSON(http.StatusConflict, gin.H{"error": "script with this path already exists"})
|
|
return
|
|
}
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to create script"})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusCreated, script)
|
|
}
|
|
|
|
// GetScript returns a script by ID.
|
|
// @Summary Get script
|
|
// @Description Returns a script by its ID
|
|
// @Tags scripts
|
|
// @Produce json
|
|
// @Param id path int true "Script ID"
|
|
// @Success 200 {object} repository.Script
|
|
// @Failure 400 {object} map[string]string
|
|
// @Failure 404 {object} map[string]string
|
|
// @Security Bearer
|
|
// @Router /scripts/:id [get]
|
|
func (sh *ScriptHandlersGroup) GetScript(c *gin.Context) {
|
|
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid id"})
|
|
return
|
|
}
|
|
|
|
script, err := sh.svc.Repo.GetScript(id)
|
|
if err != nil {
|
|
if errors.Is(err, repository.ErrNotFound) {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "script not found"})
|
|
return
|
|
}
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to get script"})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, script)
|
|
}
|
|
|
|
// UpdateScript updates a script.
|
|
// @Summary Update script
|
|
// @Description Updates a script's path, content, or interpreter
|
|
// @Tags scripts
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Param id path int true "Script ID"
|
|
// @Param body body repository.ScriptUpdate true "Script data"
|
|
// @Success 200 {object} repository.Script
|
|
// @Failure 400 {object} map[string]string
|
|
// @Failure 404 {object} map[string]string
|
|
// @Security Bearer
|
|
// @Router /scripts/:id [put]
|
|
func (sh *ScriptHandlersGroup) UpdateScript(c *gin.Context) {
|
|
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid id"})
|
|
return
|
|
}
|
|
|
|
var req repository.ScriptUpdate
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
|
|
return
|
|
}
|
|
|
|
script, err := sh.svc.Repo.UpdateScript(id, req)
|
|
if err != nil {
|
|
if errors.Is(err, repository.ErrNotFound) {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "script not found"})
|
|
return
|
|
}
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to update script"})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, script)
|
|
}
|
|
|
|
// DeleteScript deletes a script.
|
|
// @Summary Delete script
|
|
// @Description Deletes a script by its ID
|
|
// @Tags scripts
|
|
// @Param id path int true "Script ID"
|
|
// @Success 200 {object} map[string]string
|
|
// @Failure 400 {object} map[string]string
|
|
// @Failure 404 {object} map[string]string
|
|
// @Security Bearer
|
|
// @Router /scripts/:id [delete]
|
|
func (sh *ScriptHandlersGroup) DeleteScript(c *gin.Context) {
|
|
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid id"})
|
|
return
|
|
}
|
|
|
|
if err := sh.svc.Repo.DeleteScript(id); err != nil {
|
|
if errors.Is(err, repository.ErrNotFound) {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "script not found"})
|
|
return
|
|
}
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to delete script"})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{"message": "script deleted"})
|
|
}
|
|
|
|
// RunScriptByID executes a stored script on a target agent.
|
|
// @Summary Run script by ID
|
|
// @Description Loads a script from storage, resolves interpreter command, and executes on the specified agent
|
|
// @Tags scripts
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Param id path int true "Script ID"
|
|
// @Param body body RunStoredScriptIn true "Agent token and optional stdin"
|
|
// @Success 201 {object} RunScriptOut
|
|
// @Failure 400 {object} map[string]string
|
|
// @Failure 404 {object} map[string]string
|
|
// @Failure 500 {object} map[string]string
|
|
// @Security Bearer
|
|
// @Router /scripts/:id/run [post]
|
|
func (sh *ScriptHandlersGroup) RunScriptByID(c *gin.Context) {
|
|
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid id"})
|
|
return
|
|
}
|
|
|
|
var in RunStoredScriptIn
|
|
if err := c.ShouldBindJSON(&in); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
|
|
return
|
|
}
|
|
|
|
script, err := sh.svc.Repo.GetScript(id)
|
|
if err != nil {
|
|
if errors.Is(err, repository.ErrNotFound) {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "script not found"})
|
|
return
|
|
}
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to get script"})
|
|
return
|
|
}
|
|
|
|
command, err := sh.svc.ResolveCommand(c.Request.Context(), script.InterpreterID, script.Content)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to resolve command: %v", err)})
|
|
return
|
|
}
|
|
|
|
agent, ok := sh.cmder.GetAgent(in.Token)
|
|
if !ok {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "agent not found"})
|
|
return
|
|
}
|
|
|
|
jid, err := agent.AddJob(models.JobForInsert{
|
|
Command: command,
|
|
Stdin: in.Stdin,
|
|
})
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to add job: %v", err)})
|
|
return
|
|
}
|
|
|
|
job, err := agent.WaitJob(jid)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("job execution failed: %v", err)})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusCreated, RunScriptOut{
|
|
ID: job.ID,
|
|
Command: job.Command,
|
|
Stdin: job.Stdin,
|
|
Stdout: job.Stdout,
|
|
Stderr: job.Stderr,
|
|
Status: job.Status,
|
|
})
|
|
}
|
|
|
|
// RunStoredScriptIn is the request body for running a stored script on an agent.
|
|
type RunStoredScriptIn struct {
|
|
Token string `json:"token" binding:"required"`
|
|
Stdin *string `json:"stdin"`
|
|
}
|
|
|
|
// isUniqueConstraint checks if the error is a SQLite UNIQUE constraint violation.
|
|
func isUniqueConstraint(err error) bool {
|
|
return err != nil && (err.Error() != "" && contains(err.Error(), "UNIQUE constraint"))
|
|
}
|
|
|
|
func contains(s, substr string) bool {
|
|
return len(s) >= len(substr) && searchSubstring(s, substr)
|
|
}
|
|
|
|
func searchSubstring(s, substr string) bool {
|
|
for i := 0; i <= len(s)-len(substr); i++ {
|
|
if s[i:i+len(substr)] == substr {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|