Reusable fleet variable find and replace (#24613)

# Checklist for submitter

If some of the following don't apply, delete the relevant line.

<!-- Note that API documentation changes are now addressed by the
product design team. -->

- [x] Added/updated tests
- [x] Manual QA for all new/changed functionality

server/fleet/fleet_vars.go

@@ -0,0 +1,109 @@
+package fleet
+
+import (
+	"os"
+	"strings"
+)
+
+// ContainsPrefixVars scans a string for variables in the form of $VAR
+// and ${VAR} that begin with prefix, and return an array of those
+// variables with the prefix removed.
+func ContainsPrefixVars(text, prefix string) []string {
+	vars := []string{}
+	gather := func(variable string) string {
+		if strings.HasPrefix(variable, prefix) {
+			vars = append(vars, strings.TrimPrefix(variable, prefix))
+		}
+		return ""
+	}
+	os.Expand(text, gather)
+
+	return vars
+}
+
+// MaybeExpand conditionally replaces ${var} or $var in the string based on the mapping function.
+// Only repalces the variable with the mapper string if it returns true.
+// The mapper returning false will leave the original variable unchanged.
+// Based on os.Expand
+func MaybeExpand(s string, mapping func(string) (string, bool)) string {
+	var buf []byte
+	// ${} is all ASCII, so bytes are fine for this operation.
+	i := 0
+	for j := 0; j < len(s); j++ {
+		if s[j] == '$' && j+1 < len(s) {
+			if buf == nil {
+				buf = make([]byte, 0, 2*len(s))
+			}
+			buf = append(buf, s[i:j]...)
+			name, w := getShellName(s[j+1:])
+			if name == "" {
+				// Unlike the original Expand, don't
+				// eat invalid syntax, just leave it
+				// and pass over.
+				w = 0
+				buf = append(buf, s[j])
+			} else {
+				replacement, shouldReplace := mapping(name)
+				if shouldReplace {
+					buf = append(buf, replacement...)
+				} else {
+					// We aren't replacing this
+					// reference, pass over it.
+					w = 0
+					buf = append(buf, s[j])
+				}
+			}
+			j += w
+			i = j + 1
+		}
+	}
+	if buf == nil {
+		return s
+	}
+	return string(buf) + s[i:]
+}
+
+// isShellSpecialVar reports whether the character identifies a special
+// shell variable such as $*.
+func isShellSpecialVar(c uint8) bool {
+	switch c {
+	case '*', '#', '$', '@', '!', '?', '-', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
+		return true
+	}
+	return false
+}
+
+// isAlphaNum reports whether the byte is an ASCII letter, number, or underscore.
+func isAlphaNum(c uint8) bool {
+	return c == '_' || '0' <= c && c <= '9' || 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z'
+}
+
+// getShellName returns the name that begins the string and the number of bytes
+// consumed to extract it. If the name is enclosed in {}, it's part of a ${}
+// expansion and two more bytes are needed than the length of the name.
+func getShellName(s string) (string, int) {
+	switch {
+	case s[0] == '{':
+		if len(s) > 2 && isShellSpecialVar(s[1]) && s[2] == '}' {
+			return s[1:2], 3
+		}
+		// Scan to closing brace
+		for i := 1; i < len(s); i++ {
+			if s[i] == '}' {
+				if i == 1 {
+					return "", 2 // Bad syntax; eat "${}"
+				}
+				return s[1:i], i + 1
+			}
+		}
+		return "", 1 // Bad syntax; eat "${"
+	case isShellSpecialVar(s[0]):
+		return s[0:1], 1
+	}
+	// Scan alphanumerics.
+	var i int
+	// nolint:revive
+	for i = 0; i < len(s) && isAlphaNum(s[i]); i++ {
+	}
+	return s[:i], i
+}

server/fleet/fleet_vars_test.go

@@ -0,0 +1,51 @@
+package fleet
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestContainsPrefixVars(t *testing.T) {
+	script := `
+#!/bin/sh
+
+echo $FLEET_SECRET_FOO is the secret
+echo words${FLEET_SECRET_BAR}words
+$FLEET_SECRET_BAZ
+${FLEET_SECRET_QUX}
+`
+	secrets := ContainsPrefixVars(script, FLEET_SECRET_PREFIX)
+	require.Contains(t, secrets, "FOO")
+	require.Contains(t, secrets, "BAR")
+	require.Contains(t, secrets, "BAZ")
+	require.Contains(t, secrets, "QUX")
+}
+
+func TestMaybeExpand(t *testing.T) {
+	script := `
+This is $OTHER_VAR, $ $$ $* ${} in a sentence with${ALSO_OTHER_VAR}in the middle.
+We want to remember $FLEET_SECRET_BANANA and also${FLEET_SECRET_STRAWBERRY}are important.
+`
+	expected := `
+This is $OTHER_VAR, $ $$ $* ${} in a sentence with${ALSO_OTHER_VAR}in the middle.
+We want to remember BREAD and alsoSHORTCAKEare important.
+`
+
+	mapping := map[string]string{
+		"BANANA":     "BREAD",
+		"STRAWBERRY": "SHORTCAKE",
+	}
+
+	mapper := func(s string) (string, bool) {
+		if strings.HasPrefix(s, FLEET_SECRET_PREFIX) {
+			return mapping[strings.TrimPrefix(s, FLEET_SECRET_PREFIX)], true
+		}
+		return "", false
+	}
+
+	expanded := MaybeExpand(script, mapper)
+
+	require.Equal(t, expected, expanded)
+}

server/fleet/secrets.go

@@ -0,0 +1,3 @@
+package fleet
+
+const FLEET_SECRET_PREFIX = "FLEET_SECRET_"