diff --git a/agent/agent_windows.go b/agent/agent_windows.go index abdb298..f466150 100644 --- a/agent/agent_windows.go +++ b/agent/agent_windows.go @@ -48,123 +48,7 @@ var ( -func (a *Agent) 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 -} func SetDetached() *syscall.SysProcAttr { return &syscall.SysProcAttr{ diff --git a/agent/system/system_windows.go b/agent/system/system_windows.go new file mode 100644 index 0000000..f081b15 --- /dev/null +++ b/agent/system/system_windows.go @@ -0,0 +1,119 @@ +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 +} \ No newline at end of file