Implement replacing secrets in agent.download section

Author: ycombinator

internal/pkg/model/schema.go

@@ -87,6 +87,19 @@ func NewParsedPolicy(ctx context.Context, bulker bulk.Bulk, p model.Policy) (*Pa
 	policyInputs, keys := secret.ProcessInputsSecrets(p.Data, secretValues)
 	secretKeys = append(secretKeys, keys...)
 
+	// FIXME: Replace secrets in 'agent.download' section of policy
+	if agentDownload, exists := p.Data.Agent["download"]; exists {
+		if section, ok := agentDownload.(map[string]interface{}); ok {
+			agentDownloadSecretKeys, err := secret.ProcessAgentDownloadSecrets(ctx, section, bulker)
+			if err != nil {
+				return nil, fmt.Errorf("error processing agent secrets: %w", err)
+			}
+			for _, key := range agentDownloadSecretKeys {
+				secretKeys = append(secretKeys, "agent.download."+key)
+			}
+		}
+	}
+
 	// Done replacing secrets.
 	p.Data.SecretReferences = nil
 

internal/pkg/policy/parsed_policy.go

@@ -125,11 +125,12 @@ func TestParsedPolicyMixedSecretsReplacement(t *testing.T) {
 	require.NoError(t, err)
 
 	// Validate that secrets were identified
-	require.Len(t, pp.SecretKeys, 4)
+	require.Len(t, pp.SecretKeys, 5)
 	require.Contains(t, pp.SecretKeys, "outputs.fs-output.type")
 	require.Contains(t, pp.SecretKeys, "outputs.fs-output.ssl.key")
 	require.Contains(t, pp.SecretKeys, "inputs.0.streams.0.auth.basic.password")
 	require.Contains(t, pp.SecretKeys, "inputs.0.streams.1.auth.basic.password")
+	require.Contains(t, pp.SecretKeys, "agent.download.ssl.key")
 
 	// Validate that secret references were replaced
 	firstInputStreams := pp.Inputs[0]["streams"].([]any)

internal/pkg/policy/parsed_policy_test.go

@@ -274,52 +274,52 @@ func replaceInlineSecretRefsInSlice(arr []any, secrets map[string]string) ([]any
 	return result, keys
 }
 
-type OutputSecret struct {
+type Secret struct {
 	Path []string
 	ID   string
 }
 
-func getSecretIDAndPath(secret smap.Map) []OutputSecret {
-	outputSecrets := make([]OutputSecret, 0)
+func getSecretIDAndPath(secret smap.Map) []Secret {
+	secrets := make([]Secret, 0)
 
 	secretID := secret.GetString("id")
 	if secretID != "" {
-		outputSecrets = append(outputSecrets, OutputSecret{
+		secrets = append(secrets, Secret{
 			Path: make([]string, 0),
 			ID:   secretID,
 		})
 
-		return outputSecrets
+		return secrets
 	}
 
 	for secretKey := range secret {
-		newOutputSecrets := getSecretIDAndPath(secret.GetMap(secretKey))
+		newSecrets := getSecretIDAndPath(secret.GetMap(secretKey))
 
-		for _, secret := range newOutputSecrets {
-			path := append([]string{secretKey}, secret.Path...)
-			outputSecrets = append(outputSecrets, OutputSecret{
+		for _, newSecret := range newSecrets {
+			path := append([]string{secretKey}, newSecret.Path...)
+			secrets = append(secrets, Secret{
 				Path: path,
-				ID:   secret.ID,
+				ID:   newSecret.ID,
 			})
 		}
 	}
 
-	return outputSecrets
+	return secrets
 }
 
-func setSecretPath(output smap.Map, secretValue string, secretPaths []string) {
+func setSecretPath(section smap.Map, secretValue string, secretPaths []string) {
 	// Break the recursion
 	if len(secretPaths) == 1 {
-		output[secretPaths[0]] = secretValue
+		section[secretPaths[0]] = secretValue
 		return
 	}
 	path, secretPaths := secretPaths[0], secretPaths[1:]
 
-	if output.GetMap(path) == nil {
-		output[path] = make(map[string]interface{})
+	if section.GetMap(path) == nil {
+		section[path] = make(map[string]interface{})
 	}
 
-	setSecretPath(output.GetMap(path), secretValue, secretPaths)
+	setSecretPath(section.GetMap(path), secretValue, secretPaths)
 }
 
 // Read secret from output and mutate output with secret value
@@ -378,6 +378,42 @@ func processOutputWithInlineSecrets(output smap.Map, secretValues map[string]str
 	return keys
 }
 
+// ProcessAgentDownloadSecrets reads and replaces secrets in the agent.download section of the policy
+func ProcessAgentDownloadSecrets(ctx context.Context, agentDownload smap.Map, bulker bulk.Bulk) ([]string, error) {
+	secrets := agentDownload.GetMap(FieldSecrets)
+	delete(agentDownload, FieldSecrets)
+
+	secretReferences := make([]model.SecretReferencesItems, 0)
+	agentDownloadSecrets := getSecretIDAndPath(secrets)
+	keys := make([]string, 0, len(agentDownloadSecrets))
+
+	for _, secret := range agentDownloadSecrets {
+		secretReferences = append(secretReferences, model.SecretReferencesItems{
+			ID: secret.ID,
+		})
+	}
+	if len(secretReferences) == 0 {
+		return nil, nil
+	}
+	secretValues, err := GetSecretValues(ctx, secretReferences, bulker)
+	if err != nil {
+		return nil, err
+	}
+	for _, secret := range agentDownloadSecrets {
+		var key string
+		for _, p := range secret.Path {
+			if key == "" {
+				key = p
+				continue
+			}
+			key = key + "." + p
+		}
+		keys = append(keys, key)
+		setSecretPath(agentDownload, secretValues[secret.ID], secret.Path)
+	}
+	return keys, nil
+}
+
 // replaceStringRef replaces values matching a secret ref regex, e.g. $co.elastic.secret{<secret ref>} -> <secret value>
 // and does this for multiple matches
 // returns the resulting string value, and if any replacements were made

internal/pkg/secret/secret.go

@@ -171,7 +171,7 @@ func Test_Agent_Namespace_test1(t *testing.T) {
 		},
 		OutputPermissions: json.RawMessage(`{"default": {} }`),
 		Inputs:            []map[string]interface{}{},
-		Agent:             json.RawMessage(`{"monitoring": {"use_output":"default"}}`),
+		Agent:             map[string]interface{}{"monitoring": {"use_output": "default"}},
 	}
 
 	_, err = dl.CreatePolicy(ctx, srv.bulker, model.Policy{

internal/pkg/server/namespaces_integration_test.go

@@ -165,7 +165,7 @@ func Test_Agent_Remote_ES_Output(t *testing.T) {
 		},
 		OutputPermissions: json.RawMessage(`{"default": {}, "remoteES": {}}`),
 		Inputs:            []map[string]interface{}{},
-		Agent:             json.RawMessage(`{"monitoring": {"use_output":"remoteES"}}`),
+		Agent:             map[string]interface{}{"monitoring": {"use_output": "remoteES"}},
 	}
 
 	_, err = dl.CreatePolicy(ctx, srv.bulker, model.Policy{
@@ -319,7 +319,7 @@ func Test_Agent_Remote_ES_Output_ForceUnenroll(t *testing.T) {
 		},
 		OutputPermissions: json.RawMessage(`{"default": {}, "remoteES": {}}`),
 		Inputs:            []map[string]interface{}{},
-		Agent:             json.RawMessage(`{"monitoring": {"use_output":"remoteES"}}`),
+		Agent:             map[string]interface{}{"monitoring": {"use_output": "remoteES"}},
 	}
 
 	_, err = dl.CreatePolicy(ctx, srv.bulker, model.Policy{
@@ -440,7 +440,7 @@ func Test_Agent_Remote_ES_Output_Unenroll(t *testing.T) {
 		},
 		OutputPermissions: json.RawMessage(`{"default": {}, "remoteES": {}}`),
 		Inputs:            []map[string]interface{}{},
-		Agent:             json.RawMessage(`{"monitoring": {"use_output":"remoteES"}}`),
+		Agent:             map[string]interface{}{"monitoring": {"use_output": "remoteES"}},
 	}
 
 	_, err = dl.CreatePolicy(ctx, srv.bulker, model.Policy{