fleet/pkg/spec/spec_test.go

package spec

import (
	"encoding/json"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"

	"github.com/fleetdm/fleet/v4/pkg/testutils"
	"github.com/fleetdm/fleet/v4/server/fleet"
	"github.com/fleetdm/fleet/v4/server/platform/endpointer"
	"github.com/hashicorp/go-multierror"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestSplitYaml(t *testing.T) {
	in := `
---
- Document
#---
--- Document2
---
Document3
`

	docs := SplitYaml(in)
	require.Equal(t, 3, len(docs))
	assert.Equal(t, "- Document\n#---", docs[0])
	assert.Equal(t, "Document2", docs[1])
	assert.Equal(t, "Document3", docs[2])
}

func gitRootPath(t *testing.T) string {
	path, err := exec.Command("git", "rev-parse", "--show-toplevel").Output()
	require.NoError(t, err)
	return strings.TrimSpace(string(path))
}

func loadSpec(t *testing.T, relativePaths ...string) []byte {
	b, err := os.ReadFile(filepath.Join(
		append([]string{gitRootPath(t)}, relativePaths...)...,
	))
	require.NoError(t, err)
	return b
}

func TestGroupFromBytesWithStdLib(t *testing.T) {
	stdQueryLib := loadSpec(t,
		"docs", "01-Using-Fleet", "standard-query-library", "standard-query-library.yml",
	)
	g, err := GroupFromBytes(stdQueryLib)
	require.NoError(t, err)
	require.NotEmpty(t, g.Queries)
	require.NotEmpty(t, g.Policies)
}

func TestGroupFromBytesWithMacOS13CISQueries(t *testing.T) {
	cisQueries := loadSpec(t,
		"ee", "cis", "macos-13", "cis-policy-queries.yml",
	)
	g, err := GroupFromBytes(cisQueries)
	require.NoError(t, err)
	require.NotEmpty(t, g.Policies)
}

func TestGroupFromBytesWithWin10CISQueries(t *testing.T) {
	cisQueries := loadSpec(t,
		"ee", "cis", "win-10", "cis-policy-queries.yml",
	)
	g, err := GroupFromBytes(cisQueries)
	require.NoError(t, err)
	require.NotEmpty(t, g.Policies)
}

func TestGroupFromBytesMissingFields(t *testing.T) {
	tests := []struct {
		name string
		in   []byte
		want string
	}{
		{
			"missing spec",
			[]byte(`
---
apiVersion: v1
kind: team
			`),
			`Missing required fields ("spec") on provided "team" configuration.`,
		},
		{
			"missing spec and kind",
			[]byte(`
---
apiVersion: v1
			`),
			`Missing required fields ("spec", "kind") on provided configuration`,
		},
		{
			"missing spec and empty string kind",
			[]byte(`
---
apiVersion: v1
kind: ""
			`),
			`Missing required fields ("spec", "kind") on provided configuration`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, err := GroupFromBytes(tt.in)
			require.ErrorContains(t, err, tt.want)
		})
	}
}

func TestEscapeString(t *testing.T) {
	for _, tc := range []struct {
		s         string
		expResult string
	}{
		{`$foo`, `$foo`},                                      // nothing to escape
		{`bar$foo`, `bar$foo`},                                // nothing to escape
		{`bar${foo}`, `bar${foo}`},                            // nothing to escape
		{`\$foo`, `$PREVENT_ESCAPING_foo`},                    // escaping
		{`bar\$foo`, `bar$PREVENT_ESCAPING_foo`},              // escaping
		{`\\$foo`, `\\$foo`},                                  // no escaping
		{`bar\\$foo`, `bar\\$foo`},                            // no escaping
		{`\\\$foo`, `\$PREVENT_ESCAPING_foo`},                 // escaping
		{`bar\\\$foo`, `bar\$PREVENT_ESCAPING_foo`},           // escaping
		{`bar\\\${foo}bar`, `bar\$PREVENT_ESCAPING_{foo}bar`}, // escaping
		{`\\\\$foo`, `\\\\$foo`},                              // no escaping
		{`bar\\\\$foo`, `bar\\\\$foo`},                        // no escaping
		{`bar\\\\${foo}`, `bar\\\\${foo}`},                    // no escaping
	} {
		result := escapeString(tc.s, "PREVENT_ESCAPING_")
		require.Equal(t, tc.expResult, result)
	}
}

func checkMultiErrors(t *testing.T, errs ...string) func(err error) {
	return func(err error) {
		me, ok := err.(*multierror.Error)
		require.True(t, ok)
		require.Len(t, me.Errors, len(errs))
		for i, err := range me.Errors {
			require.Equal(t, errs[i], err.Error())
		}
	}
}

func TestExpandEnv(t *testing.T) {
	for _, tc := range []struct {
		environment map[string]string
		s           string
		expResult   string
		checkErr    func(error)
	}{
		{map[string]string{"foo": "1"}, `$foo`, `1`, nil},
		{map[string]string{"foo": "1"}, `$foo $FLEET_VAR_BAR ${FLEET_VAR_BAR}x ${foo}`, `1 $FLEET_VAR_BAR ${FLEET_VAR_BAR}x 1`, nil},
		{map[string]string{"foo": ""}, `$foo`, ``, nil},
		{map[string]string{"foo": "", "bar": "", "zoo": ""}, `$foo${bar}$zoo`, ``, nil},
		{map[string]string{}, `$foo`, ``, checkMultiErrors(t, "environment variable \"foo\" not set")},
		{map[string]string{"foo": "1"}, `$foo$bar`, ``, checkMultiErrors(t, "environment variable \"bar\" not set")},
		{
			map[string]string{"bar": "1"},
			`$foo $bar $zoo`, ``,
			checkMultiErrors(t, "environment variable \"foo\" not set", "environment variable \"zoo\" not set"),
		},
		{map[string]string{"foo": "4", "bar": "2"}, `$foo$bar`, `42`, nil},
		{map[string]string{"foo": "42", "bar": ""}, `$foo$bar`, `42`, nil},
		{map[string]string{}, `$$`, ``, checkMultiErrors(t, "environment variable \"$\" not set")},
		{map[string]string{"foo": "1"}, `$$foo`, ``, checkMultiErrors(t, "environment variable \"$\" not set")},
		{map[string]string{"foo": "1"}, `\$${foo}`, `$1`, nil},
		{map[string]string{}, `\$foo`, `$foo`, nil},                     // escaped
		{map[string]string{"foo": "1"}, `\\$foo`, `\\1`, nil},           // not escaped
		{map[string]string{}, `\\\$foo`, `\$foo`, nil},                  // escaped
		{map[string]string{}, `\\\$foo$`, `\$foo$`, nil},                // escaped
		{map[string]string{}, `bar\\\$foo$`, `bar\$foo$`, nil},          // escaped
		{map[string]string{"foo": "1"}, `$foo var`, `1 var`, nil},       // not escaped
		{map[string]string{"foo": "1"}, `${foo}var`, `1var`, nil},       // not escaped
		{map[string]string{"foo": "1"}, `\${foo}var`, `${foo}var`, nil}, // escaped
		{map[string]string{"foo": ""}, `${foo}var`, `var`, nil},
		{map[string]string{"foo": "", "$": "2"}, `${$}${foo}var`, `2var`, nil},
		{map[string]string{}, `${foo}var`, ``, checkMultiErrors(t, "environment variable \"foo\" not set")},
		{map[string]string{}, `foo PREVENT_ESCAPING_bar $ FLEET_VAR_`, `foo PREVENT_ESCAPING_bar $ FLEET_VAR_`, nil}, // nothing to replace
		{
			map[string]string{"foo": "BAR"},
			`\$FLEET_VAR_$foo \${FLEET_VAR_$foo} \${FLEET_VAR_${foo}2}`,
			`$FLEET_VAR_BAR ${FLEET_VAR_BAR} ${FLEET_VAR_BAR2}`, nil,
		}, // nested variables
		{
			map[string]string{},
			"$fleet_var_test", // Somehow, variables can be lowercased when coming in from time to time
			"$fleet_var_test", // Should not be replaced
			nil,
		},
		{
			map[string]string{"custom_secret": "test&123"},
			"<Add>$custom_secret</Add>",
			"<Add>test&amp;123</Add>",
			nil,
		},
	} {
		// save the current env before clearing it.
		testutils.SaveEnv(t)
		os.Clearenv()
		for k, v := range tc.environment {
			_ = os.Setenv(k, v)
		}
		result, err := ExpandEnv(tc.s)
		if tc.checkErr == nil {
			require.NoError(t, err)
		} else {
			tc.checkErr(err)
		}
		require.Equal(t, tc.expResult, result)
	}
}

func TestLookupEnvSecrets(t *testing.T) {
	for _, tc := range []struct {
		environment map[string]string
		s           string
		expResult   map[string]string
		checkErr    func(error)
	}{
		{map[string]string{"foo": "1"}, `$foo`, map[string]string{}, nil},
		{map[string]string{"FLEET_SECRET_foo": "1"}, `$FLEET_SECRET_foo`, map[string]string{"FLEET_SECRET_foo": "1"}, nil},
		{
			map[string]string{"foo": "1"},
			`$FLEET_SECRET_foo`,
			map[string]string{},
			checkMultiErrors(t, "environment variable \"FLEET_SECRET_foo\" not set"),
		},
		{map[string]string{"FLEET_SECRET_foo": "test&123"}, `<Add>$FLEET_SECRET_foo</Add>`, map[string]string{"FLEET_SECRET_foo": "test&123"}, nil},
	} {
		// save the current env before clearing it.
		testutils.SaveEnv(t)
		os.Clearenv()
		for k, v := range tc.environment {
			_ = os.Setenv(k, v)
		}
		secretsMap := make(map[string]string)
		err := LookupEnvSecrets(tc.s, secretsMap)
		if tc.checkErr == nil {
			require.NoError(t, err)
		} else {
			tc.checkErr(err)
		}
		require.Equal(t, tc.expResult, secretsMap)
	}
}

// TestExpandEnvBytesIncludingSecrets tests that FLEET_SECRET_ variables are expanded when using ExpandEnvBytesIncludingSecrets
func TestExpandEnvBytesIncludingSecrets(t *testing.T) {
	t.Setenv("FLEET_SECRET_API_KEY", "secret123")
	t.Setenv("NORMAL_VAR", "normalvalue")
	t.Setenv("FLEET_VAR_HOST", "hostname")
	t.Setenv("FLEET_SECRET_XML", "secret&123")

	input := []byte(`API Key: $FLEET_SECRET_API_KEY
Normal: $NORMAL_VAR  
Fleet Var: $FLEET_VAR_HOST
Missing: $FLEET_SECRET_MISSING`)

	result, err := ExpandEnvBytesIncludingSecrets(input)
	require.NoError(t, err)

	expected := `API Key: secret123
Normal: normalvalue  
Fleet Var: $FLEET_VAR_HOST
Missing: $FLEET_SECRET_MISSING`

	assert.Equal(t, expected, string(result))

	// Verify that FLEET_VAR_ is not expanded (reserved for server)
	assert.Contains(t, string(result), "$FLEET_VAR_HOST")
	// Verify that FLEET_SECRET_ is expanded
	assert.Contains(t, string(result), "secret123")
	assert.NotContains(t, string(result), "$FLEET_SECRET_API_KEY")
	// Verify that missing secrets are left as-is
	assert.Contains(t, string(result), "$FLEET_SECRET_MISSING")

	xmlInput := []byte(`<Add>$FLEET_SECRET_XML</Add>`)
	xmlResult, err := ExpandEnvBytesIncludingSecrets(xmlInput)
	require.NoError(t, err)

	expectedXML := `<Add>secret&amp;123</Add>`
	assert.Equal(t, expectedXML, string(xmlResult))
}

func TestGetExclusionZones(t *testing.T) {
	// Test with a small dedicated fixture where exact byte positions are stable
	t.Run("testdata/policies/policies.yml", func(t *testing.T) {
		fContents, err := os.ReadFile(filepath.Join("testdata", "policies", "policies.yml"))
		require.NoError(t, err)

		contents := string(fContents)
		actual := getExclusionZones(contents)

		expected := map[[2]int]string{
			{46, 106}:  "  description: This policy should always fail.\n  resolution:",
			{93, 155}:  "  resolution: There is no resolution for this policy.\n  query:",
			{268, 328}: "  description: This policy should always pass.\n  resolution:",
			{315, 678}: "  resolution: |\n    Automated method:\n    Ask your system administrator to deploy the following script which will ensure proper Security Auditing Retention:\n    cp /etc/security/audit_control ./tmp.txt; origExpire=$(cat ./tmp.txt  | grep expire-after);  sed \"s/${origExpire}/expire-after:60d OR 5G/\" ./tmp.txt > /etc/security/audit_control; rm ./tmp.txt;\n  query:",
		}
		require.Equal(t, len(expected), len(actual))

		for pos, text := range expected {
			assert.Contains(t, actual, pos)
			assert.Equal(t, contents[pos[0]:pos[1]], text, pos)
		}
	})

	// Test with a larger config file - verify expected text strings are found within zones
	// without hardcoding byte positions (which shift when the file is modified)
	t.Run("testdata/global_config_no_paths.yml", func(t *testing.T) {
		fContents, err := os.ReadFile(filepath.Join("testdata", "global_config_no_paths.yml"))
		require.NoError(t, err)

		contents := string(fContents)
		actual := getExclusionZones(contents)

		// Expected text strings that should be found within exclusion zones
		expectedTexts := []string{
			"    description: Collect osquery performance stats directly from osquery\n    query:",
			"    description: This policy should always fail.\n    resolution:",
			"    resolution: There is no resolution for this policy.\n    query:",
			"    description: This policy should always pass.\n    resolution:",
			"    resolution: |\n      Automated method:",
			"    description: A cool global label\n    query:",
			"    description: A fly global label\n    hosts:",
		}

		for _, expectedText := range expectedTexts {
			found := false
			for _, zone := range actual {
				zoneText := contents[zone[0]:zone[1]]
				if zoneText == expectedText || strings.Contains(zoneText, strings.TrimPrefix(expectedText, "    ")) {
					found = true
					break
				}
			}
			assert.True(t, found, "expected text not found in any exclusion zone: %q", expectedText)
		}
	})
}

func TestRewriteNewToOldKeys(t *testing.T) {
	t.Run("accepts old keys", func(t *testing.T) {
		raw := json.RawMessage(`{"name":"test","team":"myteam"}`)
		result, _, err := rewriteNewToOldKeys(raw, fleet.QuerySpec{})
		require.NoError(t, err)

		var qs fleet.QuerySpec
		require.NoError(t, json.Unmarshal(result, &qs))
		assert.Equal(t, "test", qs.Name)
		assert.Equal(t, "myteam", qs.TeamName)
	})

	t.Run("accepts new keys", func(t *testing.T) {
		raw := json.RawMessage(`{"name":"test","fleet":"myteam"}`)
		result, _, err := rewriteNewToOldKeys(raw, fleet.QuerySpec{})
		require.NoError(t, err)

		var qs fleet.QuerySpec
		require.NoError(t, json.Unmarshal(result, &qs))
		assert.Equal(t, "test", qs.Name)
		assert.Equal(t, "myteam", qs.TeamName)
	})

	t.Run("errors if both old and new keys provided", func(t *testing.T) {
		raw := json.RawMessage(`{"name":"test","team":"old","fleet":"new"}`)
		_, _, err := rewriteNewToOldKeys(raw, fleet.QuerySpec{})
		require.Error(t, err)
		var conflictErr *endpointer.AliasConflictError
		require.ErrorAs(t, err, &conflictErr)
		assert.Equal(t, "team", conflictErr.Old)
		assert.Equal(t, "fleet", conflictErr.New)
	})
}