From a5cc6a922e8221127cd8430a1b6588b3e32d1901 Mon Sep 17 00:00:00 2001
From: Paul Annesley <paul@annesley.cc>
Date: Tue, 25 Oct 2022 20:20:34 +1030
Subject: [PATCH] buildkite-agent env has --from-env-file (bool) and --print
 NAME flags

This solves the problem of safely parsing BUILDKITE_ENV_FILE, and allows
for usage like this in a pre-command hook or similar:

    local branch
    branch="$(buildkite-agent env --from-env-file --print BUILDKITE_BRANCH)"
    if [[ branch != "main" ]]; then
        echo "This agent only builds branch 'main'"
        exit 42
    fi
---
 clicommand/env.go      | 71 +++++++++++++++++++++++++++++++++++++++---
 clicommand/env_test.go | 27 ++++++++++++++++
 2 files changed, 93 insertions(+), 5 deletions(-)
 create mode 100644 clicommand/env_test.go

diff --git a/clicommand/env.go b/clicommand/env.go
index b9ab5ab5bf..62fff85507 100644
--- a/clicommand/env.go
+++ b/clicommand/env.go
@@ -1,9 +1,11 @@
 package clicommand
 
 import (
+	"bufio"
 	"encoding/json"
 	"fmt"
 	"os"
+	"strconv"
 	"strings"
 
 	"github.com/urfave/cli"
@@ -32,14 +34,33 @@ var EnvCommand = cli.Command{
 			Usage:  "Pretty print the JSON output",
 			EnvVar: "BUILDKITE_AGENT_ENV_PRETTY",
 		},
+		cli.BoolFlag{
+			Name:  "from-env-file",
+			Usage: "Source environment from file described by $BUILDKITE_ENV_FILE",
+		},
+		cli.StringFlag{
+			Name:  "print",
+			Usage: "Print a single environment variable by `NAME` as raw text followed by a newline",
+		},
 	},
 	Action: func(c *cli.Context) error {
-		env := os.Environ()
-		envMap := make(map[string]string, len(env))
+		var envMap map[string]string
+
+		if c.Bool("from-env-file") {
+			envMap = mustLoadEnvFile(os.Getenv("BUILDKITE_ENV_FILE"))
+		} else {
+			env := os.Environ()
+			envMap = make(map[string]string, len(env))
+
+			for _, e := range env {
+				k, v, _ := strings.Cut(e, "=")
+				envMap[k] = v
+			}
+		}
 
-		for _, e := range env {
-			k, v, _ := strings.Cut(e, "=")
-			envMap[k] = v
+		if name := c.String("print"); name != "" {
+			fmt.Println(envMap[name])
+			return nil
 		}
 
 		var (
@@ -69,3 +90,43 @@ var EnvCommand = cli.Command{
 		return nil
 	},
 }
+
+func mustLoadEnvFile(path string) map[string]string {
+	envMap := make(map[string]string)
+
+	if path == "" {
+		fmt.Fprintln(os.Stderr, "BUILDKITE_ENV_FILE not set")
+		os.Exit(1)
+	}
+
+	f, err := os.Open(path)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Could not open BUILDKITE_ENV_FILE: %v\n", err)
+		os.Exit(1)
+	}
+
+	scanner := bufio.NewScanner(f)
+
+	for scanner.Scan() {
+		line := scanner.Text()
+		if line == "" {
+			continue
+		}
+
+		name, quotedValue, ok := strings.Cut(line, "=")
+		if !ok {
+			fmt.Fprintf(os.Stderr, "Unexpected format in BUILDKITE_ENV_FILE %s\n", path)
+			os.Exit(1)
+		}
+
+		value, err := strconv.Unquote(quotedValue)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Error unquoting value: %v\n", err)
+			os.Exit(1)
+		}
+
+		envMap[name] = value
+	}
+
+	return envMap
+}
diff --git a/clicommand/env_test.go b/clicommand/env_test.go
new file mode 100644
index 0000000000..7359d9a8d9
--- /dev/null
+++ b/clicommand/env_test.go
@@ -0,0 +1,27 @@
+package clicommand
+
+import (
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMustLoadEnvFile(t *testing.T) {
+	f, err := os.CreateTemp("", t.Name())
+	if err != nil {
+		t.Error(err)
+	}
+	data := map[string]string{
+		"HELLO": "world",
+		"FOO":   "bar\n\"bar\"\n`black hat`\r\n$(have you any root)",
+	}
+	for name, value := range data {
+		fmt.Fprintf(f, "%s=%q\n", name, value)
+	}
+
+	result := mustLoadEnvFile(f.Name())
+
+	assert.Equal(t, data, result, "data should round-trip via env file")
+}