119 lines
No EOL
2.6 KiB
Go
119 lines
No EOL
2.6 KiB
Go
package system
|
|
|
|
func RunScript(code string, shell string, args []string, timeout int) (stdout, stderr string, exitcode int, e error) {
|
|
|
|
content := []byte(code)
|
|
|
|
dir := filepath.Join(os.TempDir(), "trmm")
|
|
if !trmm.FileExists(dir) {
|
|
a.CreateTRMMTempDir()
|
|
}
|
|
|
|
const defaultExitCode = 1
|
|
|
|
var (
|
|
outb bytes.Buffer
|
|
errb bytes.Buffer
|
|
exe string
|
|
ext string
|
|
cmdArgs []string
|
|
)
|
|
|
|
switch shell {
|
|
case "powershell":
|
|
ext = "*.ps1"
|
|
case "python":
|
|
ext = "*.py"
|
|
case "cmd":
|
|
ext = "*.bat"
|
|
}
|
|
|
|
tmpfn, err := ioutil.TempFile(dir, ext)
|
|
if err != nil {
|
|
a.Logger.Errorln(err)
|
|
return "", err.Error(), 85, err
|
|
}
|
|
defer os.Remove(tmpfn.Name())
|
|
|
|
if _, err := tmpfn.Write(content); err != nil {
|
|
a.Logger.Errorln(err)
|
|
return "", err.Error(), 85, err
|
|
}
|
|
if err := tmpfn.Close(); err != nil {
|
|
a.Logger.Errorln(err)
|
|
return "", err.Error(), 85, err
|
|
}
|
|
|
|
switch shell {
|
|
case "powershell":
|
|
exe = "Powershell"
|
|
cmdArgs = []string{"-NonInteractive", "-NoProfile", "-ExecutionPolicy", "Bypass", tmpfn.Name()}
|
|
case "python":
|
|
exe = a.PyBin
|
|
cmdArgs = []string{tmpfn.Name()}
|
|
case "cmd":
|
|
exe = tmpfn.Name()
|
|
}
|
|
|
|
if len(args) > 0 {
|
|
cmdArgs = append(cmdArgs, args...)
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
|
|
defer cancel()
|
|
|
|
var timedOut bool = false
|
|
cmd := exec.Command(exe, cmdArgs...)
|
|
cmd.Stdout = &outb
|
|
cmd.Stderr = &errb
|
|
|
|
if cmdErr := cmd.Start(); cmdErr != nil {
|
|
a.Logger.Debugln(cmdErr)
|
|
return "", cmdErr.Error(), 65, cmdErr
|
|
}
|
|
pid := int32(cmd.Process.Pid)
|
|
|
|
// custom context handling, we need to kill child procs if this is a batch script,
|
|
// otherwise it will hang forever
|
|
// the normal exec.CommandContext() doesn't work since it only kills the parent process
|
|
go func(p int32) {
|
|
|
|
<-ctx.Done()
|
|
|
|
_ = KillProc(p)
|
|
timedOut = true
|
|
}(pid)
|
|
|
|
cmdErr := cmd.Wait()
|
|
|
|
if timedOut {
|
|
stdout = CleanString(outb.String())
|
|
stderr = fmt.Sprintf("%s\nScript timed out after %d seconds", CleanString(errb.String()), timeout)
|
|
exitcode = 98
|
|
a.Logger.Debugln("Script check timeout:", ctx.Err())
|
|
} else {
|
|
stdout = CleanString(outb.String())
|
|
stderr = CleanString(errb.String())
|
|
|
|
// get the exit code
|
|
if cmdErr != nil {
|
|
if exitError, ok := cmdErr.(*exec.ExitError); ok {
|
|
if ws, ok := exitError.Sys().(syscall.WaitStatus); ok {
|
|
exitcode = ws.ExitStatus()
|
|
} else {
|
|
exitcode = defaultExitCode
|
|
}
|
|
} else {
|
|
exitcode = defaultExitCode
|
|
}
|
|
|
|
} else {
|
|
if ws, ok := cmd.ProcessState.Sys().(syscall.WaitStatus); ok {
|
|
exitcode = ws.ExitStatus()
|
|
} else {
|
|
exitcode = 0
|
|
}
|
|
}
|
|
}
|
|
return stdout, stderr, exitcode, nil
|
|
} |