src/pkg/utils/shell.go

package utils

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"os/exec"
	"regexp"
	"strings"

	"github.com/golang/glog"
)

// RunCommandString runs a command string from a provided directory and writes
// its output to the provided writer.
func RunCommandString(dir string, command string, out io.Writer) error {
	glog.Info(fmt.Sprintf("running command: %s", command))
	cmd := exec.Command("bash", "-c", command)
	cmd.Stderr = out
	cmd.Stdout = out
	cmd.Dir = dir
	return cmd.Run()
}

// SourceFile sources a bash file and returns the shell variables as a map.
func SourceFile(file string) (map[string]string, error) {
	contents, err := os.ReadFile(file)
	if err != nil {
		return nil, err
	}

	return SourceString(string(contents))
}

// SourceString sources a string as a bash file and returns the shell variables as a map.
func SourceString(contents string) (map[string]string, error) {
	// This function is a bit weird. If we were running in a shell, we could just
	// directly call `source ${path}`, and the variables set by the script ${path}
	// would be available for us. Unfortunately, that's not an option. Instead, we
	// execute contents as a script and write all of the shell variables to an
	// anonymous pipe, which we consume to populate the the mapping that we return.
	//
	// Intuitively, this is similar to executing ${path}, writing the shell
	// variables to a file, then reading them from that file, but this keeps
	// everything in RAM and we don't have to worry about filenames. The
	// anonymous pipe is passed to the child process as file descriptor 3 via
	// cmd.ExtraFiles.
	cmd := exec.Command("bash", "-c", fmt.Sprintf("%s\ndeclare -p >&3", contents))

	envReader, envWriter, err := os.Pipe()
	if err != nil {
		return nil, fmt.Errorf("failed to open pipe when sourcing bash script: %v", err)
	}
	defer envReader.Close()
	defer envWriter.Close()

	cmd.ExtraFiles = []*os.File{envWriter}
	out, err := cmd.CombinedOutput()
	glog.V(2).Info(string(out))
	if err != nil {
		return nil, fmt.Errorf("failure while sourcing bash script: %v", err)
	}
	envWriter.Close()

	scanner := bufio.NewScanner(envReader)
	env := make(map[string]string)
	for scanner.Scan() {
		// `declare -p` outputs variables like this:
		//   declare -- key="value"
		// And it outputs arrays like this:
		//   declare -A key=([0]="abc" [1]="def")
		// So we drop the first two words and take the key to be the text left of the = and
		// the value to be the text right of the =. If present, we remove one quote prefix
		// and one quote suffix from value.
		text := scanner.Text()
		parts := strings.SplitN(text, " ", 3)
		if len(parts) != 3 {
			continue
		}

		text = parts[2]
		parts = strings.SplitN(text, "=", 2)
		if len(parts) != 2 {
			continue
		}

		key := parts[0]
		value := strings.TrimSuffix(strings.TrimPrefix(parts[1], `"`), `"`)
		env[key] = value

	}
	envReader.Close()

	return env, nil
}

var arrayElementPattern = regexp.MustCompile(`\[(\d+)\]="(.*?)"`)

// ArrayElements splits a bash array output by `declare -p` into its indexed elements
// and returns the resulting map from keys to values.
// For example, the input string `X=([0]="abc" [1]="def")` should return
// map[string]string { "0": "abc", "1": "def" }.
func ArrayElements(array string) map[string]string {
	// declare escapes backslashes, which we need to un-escape
	array = strings.ReplaceAll(array, `\\`, `\`)
	// get rid of escaped quotes before we do our element matching; we make sure to keep
	// the length of all substrings the same so that we can use the indices from the
	// transformed string as the indices of matches in the original string
	escapedArray := strings.ReplaceAll(array, "\\\"", "--")
	elements := arrayElementPattern.FindAllStringSubmatchIndex(escapedArray, -1)
	result := make(map[string]string)
	for _, match := range elements {
		if len(match) != 6 {
			glog.Warningf("ill-formed array element match: %v", match)
			continue
		}
		key := array[match[2]:match[3]]
		value := array[match[4]:match[5]]
		result[key] = value
	}

	return result
}