small changed switching to windows

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{

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
+}