chore: Making the environment test cases more robust.

Signed-off-by: Matthias Glastra <matglas.git@gmail.com>

attestation/environment/environment.go

@@ -62,7 +62,7 @@ func init() {
 					return a, fmt.Errorf("unexpected attestor type: %T is not a environment attestor", a)
 				}
 
-				WithFilterVarsEnabled(filterSensitiveVarsEnabled)(envAttestor)
+				WithFilterVarsEnabled()(envAttestor)
 				return envAttestor, nil
 			},
 		),
@@ -76,7 +76,7 @@ func init() {
 					return a, fmt.Errorf("unexpected attestor type: %T is not a environment attestor", a)
 				}
 
-				WithDisableDefaultSensitiveList(disableSensitiveVarsDefault)(envAttestor)
+				WithDisableDefaultSensitiveList()(envAttestor)
 				return envAttestor, nil
 			},
 		),
@@ -103,6 +103,7 @@ type Attestor struct {
 	Username  string            `json:"username"`
 	Variables map[string]string `json:"variables,omitempty"`
 
+	osEnviron                   func() []string
 	sensitiveVarsList           map[string]struct{}
 	addSensitiveVarsList        map[string]struct{}
 	filterVarsEnabled           bool
@@ -113,9 +114,9 @@ type Option func(*Attestor)
 
 // WithFilterVarsEnabled will make the filter (removing) of vars the acting behavior.
 // The default behavior is obfuscation of variables.
-func WithFilterVarsEnabled(filterVarsEnabled bool) Option {
+func WithFilterVarsEnabled() Option {
 	return func(a *Attestor) {
-		a.filterVarsEnabled = filterVarsEnabled
+		a.filterVarsEnabled = true
 	}
 }
 
@@ -129,9 +130,16 @@ func WithAdditionalKeys(additionalKeys []string) Option {
 }
 
 // WithDisableDefaultSensitiveList will disable the default list and only use the additional keys.
-func WithDisableDefaultSensitiveList(disableSensitiveVarsDefault bool) Option {
+func WithDisableDefaultSensitiveList() Option {
 	return func(a *Attestor) {
-		a.disableSensitiveVarsDefault = disableSensitiveVarsDefault
+		a.disableSensitiveVarsDefault = true
+	}
+}
+
+// WithCustomEnv will override the default os.Environ() method. This could be used to mock.
+func WithCustomEnv(osEnviron func()[]string) Option {
+	return func(a *Attestor) {
+		a.osEnviron = osEnviron
 	}
 }
 
@@ -141,6 +149,8 @@ func New(opts ...Option) *Attestor {
 		addSensitiveVarsList: map[string]struct{}{},
 	}
 
+	attestor.osEnviron = os.Environ
+
 	for _, opt := range opts {
 		opt(attestor)
 	}
@@ -188,11 +198,11 @@ func (a *Attestor) Attest(ctx *attestation.AttestationContext) error {
 
 	// Filter or obfuscate
 	if a.filterVarsEnabled {
-		FilterEnvironmentArray(os.Environ(), finalSensitiveKeysList, func(key, val, _ string) {
+		FilterEnvironmentArray(a.osEnviron(), finalSensitiveKeysList, func(key, val, _ string) {
 			a.Variables[key] = val
 		})
 	} else {
-		ObfuscateEnvironmentArray(os.Environ(), finalSensitiveKeysList, func(key, val, _ string) {
+		ObfuscateEnvironmentArray(a.osEnviron(), finalSensitiveKeysList, func(key, val, _ string) {
 			a.Variables[key] = val
 		})
 	}

attestation/environment/environment_test.go

@@ -25,12 +25,16 @@ import (
 // TestFilterVarsEnvironment tests if enabling filter behavior works correctly.
 func TestFilterVarsEnvironment(t *testing.T) {
 
-	attestor := New(WithFilterVarsEnabled(true))
+	customEnv := func() []string {
+		return []string{"AWS_ACCESS_KEY_ID=super secret"}
+	}
+
+	attestor := New(WithFilterVarsEnabled(), WithCustomEnv(customEnv))
+
 	ctx, err := attestation.NewContext("test", []attestation.Attestor{attestor})
 	require.NoError(t, err)
 
-	t.Setenv("AWS_ACCESS_KEY_ID", "super secret")
-	origVars := os.Environ()
+	origVars := customEnv()
 	require.NoError(t, attestor.Attest(ctx))
 	for _, env := range origVars {
 		origKey, _ := splitVariable(env)
@@ -108,9 +112,42 @@ func TestEnvironmentObfuscateAdditional(t *testing.T) {
 	}
 }
 
+// TestEnvironmentCustomKeysAdditional tests if the default list is disabled the additional keys works correctly.
+func TestEnvironmentCustomKeysAdditional(t *testing.T) {
+	attestor := New(WithDisableDefaultSensitiveList(), WithAdditionalKeys([]string{"MYNAME"}))
+	ctx, err := attestation.NewContext("test", []attestation.Attestor{attestor})
+	require.NoError(t, err)
+
+	obfuscateEnvs := map[string]struct{}{"MYNAME": {}}
+	secretVarValue := "secret var"
+	publicVarValue := "public var"
+	for k := range obfuscateEnvs {
+		t.Setenv(k, secretVarValue)
+	}
+
+	notObfuscateEnvs := map[string]struct{}{"API_TOKEN": {}}
+	for k := range notObfuscateEnvs {
+		t.Setenv(k, publicVarValue)
+	}
+
+	origVars := os.Environ()
+	require.NoError(t, attestor.Attest(ctx))
+	for _, env := range origVars {
+		origKey, _ := splitVariable(env)
+		if _, inObfuscateList := obfuscateEnvs[origKey]; inObfuscateList {
+			require.NotEqual(t, attestor.Variables[origKey], secretVarValue)
+			require.Equal(t, attestor.Variables[origKey], "******")
+		}
+
+		if _, inNotObfuscateList := notObfuscateEnvs[origKey]; inNotObfuscateList {
+			require.Equal(t, attestor.Variables[origKey], publicVarValue)
+		}
+	}
+}
+
 // TestEnvironmentFilterAdditional tests if enabling filter and adding additional keys works correctly.
 func TestEnvironmentFilterAdditional(t *testing.T) {
-	attestor := New(WithFilterVarsEnabled(true), WithAdditionalKeys([]string{"MYNAME"}))
+	attestor := New(WithFilterVarsEnabled(), WithAdditionalKeys([]string{"MYNAME"}))
 	ctx, err := attestation.NewContext("test", []attestation.Attestor{attestor})
 	require.NoError(t, err)
 
@@ -139,3 +176,4 @@ func TestEnvironmentFilterAdditional(t *testing.T) {
 		}
 	}
 }
+