Updated powershell support (#18)

go.mod

@@ -4,6 +4,7 @@ go 1.17
 
 // Direct dependencies
 require (
+	github.com/alessio/shellescape v1.4.1
 	github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5
 	github.com/coreos/go-systemd/v22 v22.3.2
 	github.com/elastic/go-sysinfo v1.7.1
@@ -15,7 +16,6 @@ require (
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.9.0
 	gopkg.in/yaml.v2 v2.4.0
-	github.com/alessio/shellescape v1.4.1
 )
 
 // Transitive dependencies
@@ -50,7 +50,7 @@ require (
 	github.com/tklauser/go-sysconf v0.3.9 // indirect
 	github.com/tklauser/numcpus v0.3.0 // indirect
 	golang.org/x/crypto v0.0.0-20210921155107-089bfa567519 // indirect
-	golang.org/x/sys v0.0.0-20211015200801-69063c4bb744 // indirect
+	golang.org/x/sys v0.0.0-20211113001501-0c823b97ae02 // indirect
 	golang.org/x/text v0.3.7 // indirect
 	google.golang.org/protobuf v1.27.1 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect

sources/command/util.go

@@ -7,6 +7,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"os"
 	"os/exec"
 	"runtime"
 	"strings"
@@ -120,24 +121,19 @@ func (cp *CommandParams) Run() (*sdp.Item, error) {
 	var err error
 
 	if runtime.GOOS == "windows" {
-		return nil, &sdp.ItemRequestError{
-			ErrorType:   sdp.ItemRequestError_OTHER,
-			ErrorString: "Not currently supported on windows",
-			Context:     util.LocalContext,
-		}
+		commandString, args, err = PowerShellWrap(cp.Command, cp.Args)
 	} else {
 		commandString, args, err = ShellWrap(cp.Command, cp.Args)
+	}
 
-		if err != nil {
-			return nil, &sdp.ItemRequestError{
-				ErrorType:   sdp.ItemRequestError_OTHER,
-				ErrorString: err.Error(),
-				Context:     util.LocalContext,
-			}
+	if err != nil {
+		return nil, &sdp.ItemRequestError{
+			ErrorType:   sdp.ItemRequestError_OTHER,
+			ErrorString: err.Error(),
+			Context:     util.LocalContext,
 		}
 	}
-	// TODO: Run inside a shell
-	// TODO: Handle using powershell on windows
+
 	command := exec.CommandContext(ctx, commandString, args...)
 
 	var stdout bytes.Buffer
@@ -146,7 +142,7 @@ func (cp *CommandParams) Run() (*sdp.Item, error) {
 	command.Stdout = &stdout
 	command.Stderr = &stderr
 	command.Dir = cp.Dir
-	command.Env = envToString(cp.Env)
+	command.Env = envToString(mergeEnv(cp.Env))
 
 	err = command.Run()
 
@@ -242,6 +238,36 @@ func ShellWrap(command string, args []string) (string, []string, error) {
 	return shell, args, nil
 }
 
+// PowerShellWrap Wraps a given command and args in the required arguments so
+// that the command runs inside powershell
+func PowerShellWrap(command string, args []string) (string, []string, error) {
+	// Encode the original command as base64
+	powershellArray := []string{command}
+	powershellArray = append(powershellArray, args...)
+
+	// Append deliberate exit to ensure that the powershell.exe process exits
+	// with a code other than 1
+	powershellArray = append(powershellArray, "; exit $LASTEXITCODE")
+	powershellCommand := strings.Join(powershellArray, " ")
+
+	powershellArgs := []string{
+		"-NoLogo",          // Hides the copyright banner at startup
+		"-NoProfile",       // Does not load the Windows PowerShell profile
+		"-NonInteractive",  // Does not present an interactive prompt to the user
+		"-ExecutionPolicy", // Allow running of unsigned code
+		"bypass",
+		powershellCommand,
+	}
+
+	// TODO: The stuff that I'm doing hwere means that the actual powershell
+	// process will exit with 1 if the command fails, regardless of the exit
+	// code of the command itself. I need to find some way to pass this though.
+	// Read this:
+	// https://stackoverflow.com/questions/50200325/returning-an-exit-code-from-a-powershell-script
+
+	return "powershell.exe", powershellArgs, nil
+}
+
 // envToString Converts a map of environment variables to an array of equals
 // separated strings
 func envToString(envs map[string]string) []string {
@@ -262,3 +288,24 @@ func platformNewline() string {
 		return "\n"
 	}
 }
+
+// mergeEnv Merges the current environment variables with the supplied ones, the
+// supplied ones take precedence
+func mergeEnv(vars map[string]string) map[string]string {
+	merged := make(map[string]string)
+
+	// Get Current environment vars
+	for _, env := range os.Environ() {
+		split := strings.Split(env, "=")
+
+		if len(split) == 2 {
+			merged[split[0]] = split[1]
+		}
+	}
+
+	for k, v := range vars {
+		merged[k] = v
+	}
+
+	return merged
+}

sources/command/util_test.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"regexp"
 	"runtime"
+	"strings"
 	"testing"
 	"time"
 
@@ -363,6 +364,10 @@ func TestUnmarshalJSON(t *testing.T) {
 
 func TestShellWrap(t *testing.T) {
 	t.Run("with basic command", func(t *testing.T) {
+		if runtime.GOOS == "windows" {
+			t.Skip("skipping tests on windows")
+		}
+
 		command := "hostname"
 		args := []string{}
 
@@ -382,6 +387,10 @@ func TestShellWrap(t *testing.T) {
 	})
 
 	t.Run("with a file with spaces", func(t *testing.T) {
+		if runtime.GOOS == "windows" {
+			t.Skip("skipping tests on windows")
+		}
+
 		command := "cat"
 		args := []string{"/home/dylan/my file.txt"}
 
@@ -400,3 +409,17 @@ func TestShellWrap(t *testing.T) {
 		}
 	})
 }
+
+func TestPowershellWrap(t *testing.T) {
+	_, args, err := PowerShellWrap("Write-Host", []string{"Hello!"})
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	expected := regexp.MustCompile("Write-Host Hello!")
+
+	if !expected.MatchString(strings.Join(args, " ")) {
+		t.Fatal("Expected to match 'Write-Host Hello!'")
+	}
+}