Files
BatchuKVM/server/service/vm/script.go

156 lines
3.1 KiB
Go

package vm
import (
"fmt"
"os"
"os/exec"
"path/filepath"
"strings"
"github.com/gin-gonic/gin"
log "github.com/sirupsen/logrus"
"NanoKVM-Server/proto"
"NanoKVM-Server/utils"
)
const ScriptDirectory = "/etc/kvm/scripts"
func (s *Service) GetScripts(c *gin.Context) {
var rsp proto.Response
var files []string
err := filepath.Walk(ScriptDirectory, func(path string, info os.FileInfo, err error) error {
if err != nil {
return err
}
if !info.IsDir() && isScript(info.Name()) {
files = append(files, info.Name())
}
return nil
})
if err != nil {
rsp.ErrRsp(c, -1, "get scripts failed")
return
}
rsp.OkRspWithData(c, &proto.GetScriptsRsp{
Files: files,
})
log.Debugf("get scripts total %d", len(files))
}
func (s *Service) UploadScript(c *gin.Context) {
var rsp proto.Response
_, header, err := c.Request.FormFile("file")
if err != nil {
rsp.ErrRsp(c, -1, "bad request")
return
}
if !isScript(header.Filename) {
rsp.ErrRsp(c, -2, "invalid arguments")
return
}
if _, err = os.Stat(ScriptDirectory); err != nil {
_ = os.MkdirAll(ScriptDirectory, 0o755)
}
target := fmt.Sprintf("%s/%s", ScriptDirectory, header.Filename)
err = c.SaveUploadedFile(header, target)
if err != nil {
rsp.ErrRsp(c, -2, "save failed")
return
}
_ = utils.EnsurePermission(target, 0o100)
data := &proto.UploadScriptRsp{
File: header.Filename,
}
rsp.OkRspWithData(c, data)
log.Debugf("upload script %s success", header.Filename)
}
func (s *Service) RunScript(c *gin.Context) {
var req proto.RunScriptReq
var rsp proto.Response
if err := proto.ParseFormRequest(c, &req); err != nil {
rsp.ErrRsp(c, -1, "invalid arguments")
return
}
command := fmt.Sprintf("%s/%s", ScriptDirectory, req.Name)
name := strings.ToLower(req.Name)
if strings.HasSuffix(name, ".py") {
command = fmt.Sprintf("python %s", command)
}
var output []byte
var err error
cmd := exec.Command("sh", "-c", command)
if req.Type == "foreground" {
output, err = cmd.CombinedOutput()
} else {
cmd.Stdout = nil
cmd.Stderr = nil
go func() {
err := cmd.Run()
if err != nil {
log.Errorf("run script %s in background failed: %s", req.Name, err)
}
}()
}
if err != nil {
log.Errorf("run script %s failed: %s", req.Name, err.Error())
rsp.ErrRsp(c, -2, "run script failed")
return
}
rsp.OkRspWithData(c, &proto.RunScriptRsp{
Log: string(output),
})
log.Debugf("run script %s success", req.Name)
}
func (s *Service) DeleteScript(c *gin.Context) {
var req proto.DeleteScriptReq
var rsp proto.Response
if err := proto.ParseFormRequest(c, &req); err != nil {
rsp.ErrRsp(c, -1, "invalid arguments")
return
}
file := fmt.Sprintf("%s/%s", ScriptDirectory, req.Name)
if err := os.Remove(file); err != nil {
log.Errorf("delete script %s failed: %s", file, err)
rsp.ErrRsp(c, -3, "delete failed")
return
}
rsp.OkRsp(c)
log.Debugf("delete script %s success", file)
}
func isScript(name string) bool {
nameLower := strings.ToLower(name)
if strings.HasSuffix(nameLower, ".sh") || strings.HasSuffix(nameLower, ".py") {
return true
}
return false
}