Add support for running powershell scripts and commands.

Add test files new powershell commands feature.

Signed-off-by: Aliwoto <[email protected]>

Author: Aliwoto

ssg/helpers.go

@@ -260,10 +260,20 @@ func RunCommandAsync(command string) *ExecuteCommandResult {
 	return shellUtils.RunCommandAsync(command)
 }
 
+// RunCommandAsyncWithChan runs the given command in another goroutine and returns back
+// from the function immediately. It will sends `true` to `finishedChan` the moment
+// it's done executing the command.
 func RunCommandAsyncWithChan(command string, finishedChan chan bool) *ExecuteCommandResult {
 	return shellUtils.RunCommandAsyncWithChan(command, finishedChan)
 }
 
+// RunCommandAsyncWithChan runs the given command in another goroutine and returns back
+// from the function immediately. It will sends `true` to `finishedChan` the moment
+// it's done executing the command.
+func RunPowerShellAsyncWithChan(command string, finishedChan chan bool) *ExecuteCommandResult {
+	return shellUtils.RunPowerShellAsyncWithChan(command, finishedChan)
+}
+
 func ToBool(str string) bool {
 	return strongParser.BoolMapping[strings.ToLower(strings.TrimSpace(str))]
 }

ssg/shellUtils/constants.go

@@ -7,4 +7,7 @@ const (
 		" git rev-parse --verify HEAD &&" +
 		" git fetch && " +
 		"git rev-list --left-right --count origin/master...master"
+	PowerShellCmd            = "powershell"
+	PromptIgnoreSSG          = "PROMPT_IGNORE_SSG>"
+	PowerShellPromptOverride = "function prompt { \"" + PromptIgnoreSSG + "\" }"
 )

ssg/shellUtils/helpers.go

@@ -2,8 +2,10 @@ package shellUtils
 
 import (
 	"bytes"
+	"fmt"
 	"os"
 	"os/exec"
+	"strings"
 	"sync"
 )
 
@@ -19,22 +21,63 @@ func RunCommandAsyncWithChan(command string, finishedChan chan bool) *ExecuteCom
 	return runCommand(command, true, finishedChan)
 }
 
+func RunPowerShellAsyncWithChan(command string, finishedChan chan bool) *ExecuteCommandResult {
+	return runPowerShell(command, true, finishedChan)
+}
+
 func runCommand(command string, isAsync bool, finishedChan chan bool) *ExecuteCommandResult {
 	stdout := new(bytes.Buffer)
 	stderr := new(bytes.Buffer)
 	var result *ExecuteCommandResult
 	if isAsync {
 		result = executeCommand(command, &ExecuteCommandConfig{
+			TargetRunner:  GetCommandTargetRunner(),
+			PrimaryArgs:   GetCommandPrimaryArgs(),
 			Stdout:        stdout,
 			Stderr:        stderr,
 			autoSetOutput: true,
 			FinishedChan:  finishedChan,
-		}, true)
+			IsAsync:       true,
+		})
+		return result
+	} else {
+		result = executeCommand(command, &ExecuteCommandConfig{
+			TargetRunner: GetCommandTargetRunner(),
+			PrimaryArgs:  GetCommandPrimaryArgs(),
+			Stdout:       stdout,
+			Stderr:       stderr,
+		})
+	}
+
+	result.Stdout = stdout.String()
+	result.Stderr = stderr.String()
+
+	return result
+}
+
+func runPowerShell(command string, isAsync bool, finishedChan chan bool) *ExecuteCommandResult {
+	stdout := new(bytes.Buffer)
+	stderr := new(bytes.Buffer)
+	var result *ExecuteCommandResult
+	if isAsync {
+		result = executePowerShell(command, &ExecuteCommandConfig{
+			TargetRunner:           GetPowerShellRunner(),
+			PrimaryArgs:            GetPowerShellPrimaryArgs(),
+			Stdout:                 stdout,
+			Stderr:                 stderr,
+			autoSetOutput:          true,
+			FinishedChan:           finishedChan,
+			IsAsync:                true,
+			RemovePowerShellPrompt: true,
+		})
 		return result
 	} else {
-		result = ExecuteCommand(command, &ExecuteCommandConfig{
-			Stdout: stdout,
-			Stderr: stderr,
+		result = executePowerShell(command, &ExecuteCommandConfig{
+			TargetRunner:           GetCommandTargetRunner(),
+			PrimaryArgs:            GetPowerShellPrimaryArgs(),
+			Stdout:                 stdout,
+			Stderr:                 stderr,
+			RemovePowerShellPrompt: true,
 		})
 	}
 
@@ -45,45 +88,165 @@ func runCommand(command string, isAsync bool, finishedChan chan bool) *ExecuteCo
 }
 
 func ExecuteCommand(command string, config *ExecuteCommandConfig) *ExecuteCommandResult {
-	return executeCommand(command, config, false)
+	if config == nil {
+		config = &ExecuteCommandConfig{
+			TargetRunner: GetCommandTargetRunner(),
+			PrimaryArgs:  GetCommandPrimaryArgs(),
+		}
+	} else if config.IsAsync {
+		config.IsAsync = false
+	}
+
+	return executeCommand(command, config)
 }
 
 func ExecuteCommandAsync(command string, config *ExecuteCommandConfig) *ExecuteCommandResult {
-	return executeCommand(command, config, true)
+	if config == nil {
+		config = &ExecuteCommandConfig{
+			TargetRunner: GetCommandTargetRunner(),
+			PrimaryArgs:  GetCommandPrimaryArgs(),
+			IsAsync:      true,
+		}
+	}
+
+	return executeCommand(command, config)
+}
+
+func ExecutePowerShellAsync(command string, config *ExecuteCommandConfig) *ExecuteCommandResult {
+	if config == nil {
+		config = &ExecuteCommandConfig{
+			TargetRunner: GetPowerShellRunner(),
+			PrimaryArgs:  GetPowerShellPrimaryArgs(),
+			IsAsync:      true,
+		}
+	}
+
+	return executePowerShell(command, config)
+}
+
+func GetCommandTargetRunner() string {
+	if os.PathSeparator == '/' {
+		return ShellToUseUnix
+	} else {
+		return ShellToUseWin
+	}
+}
+
+func GetPowerShellRunner() string {
+	return PowerShellCmd
+}
+
+func GetCommandPrimaryArgs() []string {
+	if os.PathSeparator == '/' {
+		return []string{"-c"}
+	}
+	return []string{"/c"}
 }
 
+func GetPowerShellPrimaryArgs() []string {
+	return []string{"-nologo", "-noprofile", "-NonInteractive"}
+}
+
+// executeCommand is the internal version of the execute command function.
+// WARNING: the config argument MUST NOT be nil.
 func executeCommand(
 	command string,
 	config *ExecuteCommandConfig,
-	isAsync bool,
 ) *ExecuteCommandResult {
-	if config == nil {
-		config = &ExecuteCommandConfig{}
+	var cmd *exec.Cmd
+	result := &ExecuteCommandResult{
+		autoSetOutput: config.autoSetOutput,
+		mutex:         &sync.Mutex{},
+	}
+
+	cmd = exec.Command(config.TargetRunner, append(config.PrimaryArgs, command)...)
+	if len(config.ExtraFiles) != 0 {
+		cmd.ExtraFiles = append(cmd.ExtraFiles, config.ExtraFiles...)
 	}
 
+	cmd.Stdout = config.Stdout
+	cmd.Stdin = config.Stdin
+	cmd.Stderr = config.Stderr
+	cmd.Args = append(cmd.Args, config.AdditionalArgs...)
+	cmd.Env = append(cmd.Env, config.AdditionalEnv...)
+
+	result.cmd = cmd
+	result.FinishedChan = config.FinishedChan
+
+	finishUpCommand(cmd, config, result)
+
+	return result
+}
+
+// executeCommand is the internal version of the execute powershell function.
+// executes the given powershell script/command/set of commands using the
+// "powershell" (it might be powershell 5.1 which ships with windows by default).
+// WARNING: the config argument MUST NOT be nil.
+func executePowerShell(
+	command string,
+	config *ExecuteCommandConfig,
+) *ExecuteCommandResult {
 	var cmd *exec.Cmd
 	result := &ExecuteCommandResult{
 		autoSetOutput: config.autoSetOutput,
 		mutex:         &sync.Mutex{},
 	}
-	if os.PathSeparator == '/' {
-		cmd = exec.Command(ShellToUseUnix, "-c", command)
+
+	if config.RemovePowerShellPrompt && !strings.Contains(command, "function prompt") {
+		// hacky way of getting rid of powershell prompt
+		command = PowerShellPromptOverride + command
+	}
+
+	cmd = exec.Command(config.TargetRunner, config.PrimaryArgs...)
+	if len(config.ExtraFiles) != 0 {
 		cmd.ExtraFiles = append(cmd.ExtraFiles, config.ExtraFiles...)
-	} else {
-		cmd = exec.Command(ShellToUseWin, "/C", command)
 	}
 
 	cmd.Stdout = config.Stdout
-	cmd.Stdin = config.Stdin
+
+	pStdin, err := cmd.StdinPipe()
+	if err != nil {
+		result.Error = err
+		return result
+	}
+
+	if config.Stdin != nil {
+		wrappedStdin := &StdinWrapper{
+			InnerWriter: pStdin,
+		}
+
+		result.pipedStdin = wrappedStdin
+		wrappedStdin.OnWrite = append(wrappedStdin.OnWrite, func(p []byte) (n int, err error) {
+			return config.Stdin.Read(p)
+		})
+	} else {
+		result.pipedStdin = pStdin
+	}
+
 	cmd.Stderr = config.Stderr
 	cmd.Args = append(cmd.Args, config.AdditionalArgs...)
 	cmd.Env = append(cmd.Env, config.AdditionalEnv...)
 
 	result.cmd = cmd
 	result.FinishedChan = config.FinishedChan
+	_, err = fmt.Fprint(result.pipedStdin, command)
+	if err != nil {
+		result.Error = err
+		return result
+	}
 
-	if isAsync {
+	finishUpCommand(cmd, config, result)
+	return result
+}
+
+func finishUpCommand(
+	cmd *exec.Cmd,
+	config *ExecuteCommandConfig,
+	result *ExecuteCommandResult,
+) {
+	if config.IsAsync {
 		go func() {
+			result.ClosePipes()
 			result.Error = cmd.Run()
 			result.IsFinished = true
 			if result.autoSetOutput {
@@ -103,11 +266,11 @@ func executeCommand(
 			}
 		}()
 	} else {
+		result.ClosePipes()
 		result.Error = cmd.Run()
 		result.IsFinished = true
 	}
 
-	return result
 }
 
 // GetGitStats function will return the git stats in the following format:

ssg/shellUtils/methods.go

@@ -1,6 +1,7 @@
 package shellUtils
 
 import (
+	"strings"
 	"time"
 )
 
@@ -85,3 +86,85 @@ func (r *ExecuteCommandResult) Release() error {
 
 	return r.cmd.Process.Release()
 }
+
+func (r *ExecuteCommandResult) ClosePipes() error {
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
+
+	if r.pipedStdin != nil {
+		err := r.pipedStdin.Close()
+		r.pipedStdin = nil
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// PurifyOutput will purify the stdout of the result and will
+// set the purified version of the output. If you just want to
+// get the purified version without setting the new value, use
+func (r *ExecuteCommandResult) PurifyPowerShellOutput() string {
+	r.Stdout = r.GetPurifiedPowerShellOutput()
+	return r.Stdout
+}
+
+func (r *ExecuteCommandResult) GetPurifiedPowerShellOutput() string {
+	myStrs := strings.Split(r.Stdout, "\n")
+	if len(myStrs) == 0 {
+		return ""
+	}
+
+	// if the override doesn't exist, don't do any operation, just
+	// return the original value.
+	if !strings.Contains(myStrs[0], PowerShellPromptOverride) {
+		return r.Stdout
+	}
+
+	purifiedOutput := &strings.Builder{}
+	for i, currentLine := range myStrs {
+		if i == 0 {
+			// it's guaranteed that the first line is our prompt override.
+			continue
+		}
+
+		if strings.HasPrefix(currentLine, PromptIgnoreSSG) {
+			continue
+		}
+
+		purifiedOutput.WriteString(currentLine)
+		purifiedOutput.WriteRune('\n')
+	}
+
+	return strings.TrimSpace(purifiedOutput.String())
+}
+
+// --------------------------------------------------------
+
+func (r *StdinWrapper) Write(p []byte) (n int, err error) {
+	if len(r.OnWrite) > 0 {
+		for _, currentHandler := range r.OnWrite {
+			n, err = currentHandler(p)
+			if err != nil {
+				return n, err
+			}
+		}
+	}
+
+	return r.InnerWriter.Write(p)
+}
+
+func (r *StdinWrapper) Close() error {
+	if len(r.OnClose) > 0 {
+		var err error
+		for _, currentHandler := range r.OnClose {
+			err = currentHandler()
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return r.InnerWriter.Close()
+}