Implement replacing secrets in agent.download section

Author: ycombinator

internal/pkg/model/schema.go

@@ -86,6 +86,19 @@ func NewParsedPolicy(ctx context.Context, bulker bulk.Bulk, p model.Policy) (*Pa
 	}
 	secretKeys = append(secretKeys, keys...)
 
+	// 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)
+			}
+		}
+	}
+
 	// We are cool and the gang
 	pp := &ParsedPolicy{
 		Policy:  p,

internal/pkg/policy/parsed_policy.go

@@ -17,7 +17,7 @@ import (
 )
 
 const (
-	FieldOutputSecrets = "secrets"
+	FieldSecrets = "secrets"
 )
 
 var (
@@ -165,63 +165,63 @@ func replaceSliceRef(arr []any, secrets map[string]string) ([]any, []string) {
 	return result, keys
 }
 
-type OutputSecret struct {
+type Secret struct {
 	Path []string
 	ID   string
 }
 
-func getSecretIDAndPath(secret smap.Map) ([]OutputSecret, error) {
-	outputSecrets := make([]OutputSecret, 0)
+func getSecretIDAndPath(secret smap.Map) ([]Secret, error) {
+	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, nil
+		return secrets, nil
 	}
 
 	for secretKey := range secret {
-		newOutputSecrets, err := getSecretIDAndPath(secret.GetMap(secretKey))
+		newSecrets, err := getSecretIDAndPath(secret.GetMap(secretKey))
 		if err != nil {
 			return nil, err
 		}
 
-		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, nil
+	return secrets, nil
 }
 
-func setSecretPath(output smap.Map, secretValue string, secretPaths []string) error {
+func setSecretPath(section smap.Map, secretValue string, secretPaths []string) error {
 	// Break the recursion
 	if len(secretPaths) == 1 {
-		output[secretPaths[0]] = secretValue
+		section[secretPaths[0]] = secretValue
 
 		return nil
 	}
 	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{})
 	}
 
-	return setSecretPath(output.GetMap(path), secretValue, secretPaths)
+	return setSecretPath(section.GetMap(path), secretValue, secretPaths)
 }
 
 // Read secret from output and mutate output with secret value
 func ProcessOutputSecret(ctx context.Context, output smap.Map, bulker bulk.Bulk) ([]string, error) {
-	secrets := output.GetMap(FieldOutputSecrets)
+	secrets := output.GetMap(FieldSecrets)
 
-	delete(output, FieldOutputSecrets)
+	delete(output, FieldSecrets)
 	secretReferences := make([]model.SecretReferencesItems, 0)
 	outputSecrets, err := getSecretIDAndPath(secrets)
 	keys := make([]string, 0, len(outputSecrets))
@@ -259,6 +259,48 @@ func ProcessOutputSecret(ctx context.Context, output smap.Map, bulker bulk.Bulk)
 	return keys, nil
 }
 
+// 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, err := getSecretIDAndPath(secrets)
+	keys := make([]string, 0, len(agentDownloadSecrets))
+	if err != nil {
+		return nil, err
+	}
+
+	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)
+		err = setSecretPath(agentDownload, secretValues[secret.ID], secret.Path)
+		if err != nil {
+			return nil, err
+		}
+	}
+	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{