Merge pull request #1505 from Infisical/daniel/agent-improvements

Feat: Agent exec and custom polling interval

cli/agent-config.yaml

@@ -1,5 +1,5 @@
 infisical:
-  address: "http://localhost:8080"
+  address: "https://app.infisical.com/"
 auth:
   type: "universal-auth"
   config:
@@ -13,3 +13,12 @@ sinks:
 templates:
   - source-path: my-dot-ev-secret-template
     destination-path: my-dot-env.env
+    config:
+      polling-interval: 60s
+      execute:
+        command: docker-compose -f docker-compose.prod.yml down && docker-compose -f docker-compose.prod.yml up -d
+  - source-path: my-dot-ev-secret-template1
+    destination-path: my-dot-env-1.env
+    config:
+      exec:
+        command: mkdir hello-world1

cli/packages/api/api.go

@@ -490,5 +490,7 @@ func CallGetRawSecretsV3(httpClient *resty.Client, request GetRawSecretsV3Reques
 		return GetRawSecretsV3Response{}, fmt.Errorf("CallGetRawSecretsV3: Unsuccessful response [%v %v] [status-code=%v] [response=%v]", response.Request.Method, response.Request.URL, response.StatusCode(), response.String())
 	}
 
+	getRawSecretsV3Response.ETag = response.Header().Get(("etag"))
+
 	return getRawSecretsV3Response, nil
 }

cli/packages/api/model.go

@@ -505,4 +505,5 @@ type GetRawSecretsV3Response struct {
 		SecretComment string `json:"secretComment"`
 	} `json:"secrets"`
 	Imports []any `json:"imports"`
+	ETag		string
 }

cli/packages/cmd/agent.go

@@ -5,12 +5,15 @@ package cmd
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	"fmt"
 	"io/ioutil"
 	"os"
+	"os/exec"
 	"os/signal"
 	"path"
+	"runtime"
 	"strings"
 	"sync"
 	"syscall"
@@ -71,12 +74,56 @@ type Template struct {
 	SourcePath            string `yaml:"source-path"`
 	Base64TemplateContent string `yaml:"base64-template-content"`
 	DestinationPath       string `yaml:"destination-path"`
+
+	Config struct { // Configurations for the template
+		PollingInterval string `yaml:"polling-interval"` // How often to poll for changes in the secret
+		Execute         struct {
+			Command string `yaml:"command"` // Command to execute once the template has been rendered
+			Timeout int64  `yaml:"timeout"` // Timeout for the command
+		} `yaml:"execute"` // Command to execute once the template has been rendered
+	} `yaml:"config"`
 }
 
 func ReadFile(filePath string) ([]byte, error) {
 	return ioutil.ReadFile(filePath)
 }
 
+func ExecuteCommandWithTimeout(command string, timeout int64) error {
+
+	shell := [2]string{"sh", "-c"}
+	if runtime.GOOS == "windows" {
+		shell = [2]string{"cmd", "/C"}
+	} else {
+		currentShell := os.Getenv("SHELL")
+		if currentShell != "" {
+			shell[0] = currentShell
+		}
+	}
+
+	ctx := context.Background()
+	if timeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
+		defer cancel()
+	}
+
+	cmd := exec.CommandContext(ctx, shell[0], shell[1], command)
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	if err := cmd.Run(); err != nil {
+		if exitError, ok := err.(*exec.ExitError); ok { // type assertion
+			if exitError.ProcessState.ExitCode() == -1 {
+				return fmt.Errorf("command timed out")
+			}
+		}
+		return err
+	} else {
+		return nil
+	}
+}
+
 func FileExists(filepath string) bool {
 	info, err := os.Stat(filepath)
 	if os.IsNotExist(err) {
@@ -170,20 +217,24 @@ func ParseAgentConfig(configFile []byte) (*Config, error) {
 	return config, nil
 }
 
-func secretTemplateFunction(accessToken string) func(string, string, string) ([]models.SingleEnvironmentVariable, error) {
+func secretTemplateFunction(accessToken string, existingEtag string, currentEtag *string) func(string, string, string) ([]models.SingleEnvironmentVariable, error) {
 	return func(projectID, envSlug, secretPath string) ([]models.SingleEnvironmentVariable, error) {
-		secrets, err := util.GetPlainTextSecretsViaMachineIdentity(accessToken, projectID, envSlug, secretPath, false)
+		res, err := util.GetPlainTextSecretsViaMachineIdentity(accessToken, projectID, envSlug, secretPath, false)
 		if err != nil {
 			return nil, err
 		}
 
-		return secrets, nil
+		if existingEtag != res.Etag {
+			*currentEtag = res.Etag
+		}
+
+		return res.Secrets, nil
 	}
 }
 
-func ProcessTemplate(templatePath string, data interface{}, accessToken string) (*bytes.Buffer, error) {
+func ProcessTemplate(templatePath string, data interface{}, accessToken string, existingEtag string, currentEtag *string) (*bytes.Buffer, error) {
 	// custom template function to fetch secrets from Infisical
-	secretFunction := secretTemplateFunction(accessToken)
+	secretFunction := secretTemplateFunction(accessToken, existingEtag, currentEtag)
 	funcs := template.FuncMap{
 		"secret": secretFunction,
 	}
@@ -203,7 +254,7 @@ func ProcessTemplate(templatePath string, data interface{}, accessToken string)
 	return &buf, nil
 }
 
-func ProcessBase64Template(encodedTemplate string, data interface{}, accessToken string) (*bytes.Buffer, error) {
+func ProcessBase64Template(encodedTemplate string, data interface{}, accessToken string, existingEtag string, currentEtag *string) (*bytes.Buffer, error) {
 	// custom template function to fetch secrets from Infisical
 	decoded, err := base64.StdEncoding.DecodeString(encodedTemplate)
 	if err != nil {
@@ -212,7 +263,7 @@ func ProcessBase64Template(encodedTemplate string, data interface{}, accessToken
 
 	templateString := string(decoded)
 
-	secretFunction := secretTemplateFunction(accessToken)
+	secretFunction := secretTemplateFunction(accessToken, existingEtag, currentEtag) // TODO: Fix this
 	funcs := template.FuncMap{
 		"secret": secretFunction,
 	}
@@ -250,7 +301,16 @@ type TokenManager struct {
 }
 
 func NewTokenManager(fileDeposits []Sink, templates []Template, clientIdPath string, clientSecretPath string, newAccessTokenNotificationChan chan bool, removeClientSecretOnRead bool, exitAfterAuth bool) *TokenManager {
-	return &TokenManager{filePaths: fileDeposits, templates: templates, clientIdPath: clientIdPath, clientSecretPath: clientSecretPath, newAccessTokenNotificationChan: newAccessTokenNotificationChan, removeClientSecretOnRead: removeClientSecretOnRead, exitAfterAuth: exitAfterAuth}
+	return &TokenManager{
+		filePaths:                      fileDeposits,
+		templates:                      templates,
+		clientIdPath:                   clientIdPath,
+		clientSecretPath:               clientSecretPath,
+		newAccessTokenNotificationChan: newAccessTokenNotificationChan,
+		removeClientSecretOnRead:       removeClientSecretOnRead,
+		exitAfterAuth:                  exitAfterAuth,
+	}
+
 }
 
 func (tm *TokenManager) SetToken(token string, accessTokenTTL time.Duration, accessTokenMaxTTL time.Duration) {
@@ -428,38 +488,80 @@ func (tm *TokenManager) WriteTokenToFiles() {
 	}
 }
 
-func (tm *TokenManager) FetchSecrets() {
-	log.Info().Msgf("template engine started...")
+func (tm *TokenManager) WriteTemplateToFile(bytes *bytes.Buffer, template *Template) {
+	if err := WriteBytesToFile(bytes, template.DestinationPath); err != nil {
+		log.Error().Msgf("template engine: unable to write secrets to path because %s. Will try again on next cycle", err)
+		return
+	}
+	log.Info().Msgf("template engine: secret template at path %s has been rendered and saved to path %s", template.SourcePath, template.DestinationPath)
+}
+
+func (tm *TokenManager) MonitorSecretChanges(secretTemplate Template, sigChan chan os.Signal) {
+
+	pollingInterval := time.Duration(5 * time.Minute)
+
+	if secretTemplate.Config.PollingInterval != "" {
+		interval, err := util.ConvertPollingIntervalToTime(secretTemplate.Config.PollingInterval)
+
+		if err != nil {
+			log.Error().Msgf("unable to convert polling interval to time because %v", err)
+			sigChan <- syscall.SIGINT
+			return
+
+		} else {
+			pollingInterval = interval
+		}
+	}
+
+	var existingEtag string
+	var currentEtag string
+	var firstRun = true
+
+	execTimeout := secretTemplate.Config.Execute.Timeout
+	execCommand := secretTemplate.Config.Execute.Command
+
 	for {
 		token := tm.GetToken()
+
 		if token != "" {
-			for _, secretTemplate := range tm.templates {
-				var processedTemplate *bytes.Buffer
-				var err error
-				if secretTemplate.SourcePath != "" {
-					processedTemplate, err = ProcessTemplate(secretTemplate.SourcePath, nil, token)
-				} else {
-					processedTemplate, err = ProcessBase64Template(secretTemplate.Base64TemplateContent, nil, token)
-				}
 
-				if err != nil {
-					log.Error().Msgf("template engine: unable to render secrets because %s. Will try again on next cycle", err)
+			var processedTemplate *bytes.Buffer
+			var err error
 
-					continue
-				}
-
-				if err := WriteBytesToFile(processedTemplate, secretTemplate.DestinationPath); err != nil {
-					log.Error().Msgf("template engine: unable to write secrets to path because %s. Will try again on next cycle", err)
-
-					continue
-				}
-
-				log.Info().Msgf("template engine: secret template at path %s has been rendered and saved to path %s", secretTemplate.SourcePath, secretTemplate.DestinationPath)
+			if secretTemplate.SourcePath != "" {
+				processedTemplate, err = ProcessTemplate(secretTemplate.SourcePath, nil, token, existingEtag, &currentEtag)
+			} else {
+				processedTemplate, err = ProcessBase64Template(secretTemplate.Base64TemplateContent, nil, token, existingEtag, &currentEtag)
 			}
 
-			// fetch new secrets every 5 minutes (TODO: add PubSub in the future )
-			time.Sleep(5 * time.Minute)
+			if err != nil {
+				log.Error().Msgf("unable to process template because %v", err)
+			} else {
+				if (existingEtag != currentEtag) || firstRun {
+
+					tm.WriteTemplateToFile(processedTemplate, &secretTemplate)
+					existingEtag = currentEtag
+
+					if !firstRun && execCommand != "" {
+						log.Info().Msgf("executing command: %s", execCommand)
+						err := ExecuteCommandWithTimeout(execCommand, execTimeout)
+
+						if err != nil {
+							log.Error().Msgf("unable to execute command because %v", err)
+						}
+
+					}
+					if firstRun {
+						firstRun = false
+					}
+				}
+			}
+			time.Sleep(pollingInterval)
+		} else {
+			// It fails to get the access token. So we will re-try in 3 seconds. We do this because if we don't, the user will have to wait for the next polling interval to get the first secret render.
+			time.Sleep(3 * time.Second)
 		}
+
 	}
 }
 
@@ -544,7 +646,11 @@ var agentCmd = &cobra.Command{
 		tm := NewTokenManager(filePaths, agentConfig.Templates, configUniversalAuthType.ClientIDPath, configUniversalAuthType.ClientSecretPath, tokenRefreshNotifier, configUniversalAuthType.RemoveClientSecretOnRead, agentConfig.Infisical.ExitAfterAuth)
 
 		go tm.ManageTokenLifecycle()
-		go tm.FetchSecrets()
+
+		for i, template := range agentConfig.Templates {
+			log.Info().Msgf("template engine started for template %v...", i+1)
+			go tm.MonitorSecretChanges(template, sigChan)
+		}
 
 		for {
 			select {

cli/packages/models/cli.go

@@ -34,6 +34,11 @@ type SingleEnvironmentVariable struct {
 	Comment string `json:"comment"`
 }
 
+type PlaintextSecretResult struct {
+	Secrets []SingleEnvironmentVariable
+	Etag    string
+}
+
 type SingleFolder struct {
 	ID   string `json:"_id"`
 	Name string `json:"name"`

cli/packages/util/agent.go

@@ -0,0 +1,41 @@
+package util
+
+import (
+	"fmt"
+	"strconv"
+	"time"
+)
+
+// ConvertPollingIntervalToTime converts a string representation of a polling interval to a time.Duration
+func ConvertPollingIntervalToTime(pollingInterval string) (time.Duration, error) {
+	length := len(pollingInterval)
+	if length < 2 {
+		return 0, fmt.Errorf("invalid format")
+	}
+
+	unit := pollingInterval[length-1:]
+	numberPart := pollingInterval[:length-1]
+
+	number, err := strconv.Atoi(numberPart)
+	if err != nil {
+		return 0, err
+	}
+
+	switch unit {
+	case "s":
+		if number < 60 {
+			return 0, fmt.Errorf("polling interval should be at least 60 seconds")
+		}
+		return time.Duration(number) * time.Second, nil
+	case "m":
+		return time.Duration(number) * time.Minute, nil
+	case "h":
+		return time.Duration(number) * time.Hour, nil
+	case "d":
+		return time.Duration(number) * 24 * time.Hour, nil
+	case "w":
+		return time.Duration(number) * 7 * 24 * time.Hour, nil
+	default:
+		return 0, fmt.Errorf("invalid time unit")
+	}
+}

cli/packages/util/secrets.go

@@ -152,7 +152,7 @@ func GetPlainTextSecretsViaJTW(JTWToken string, receiversPrivateKey string, work
 	return plainTextSecrets, nil
 }
 
-func GetPlainTextSecretsViaMachineIdentity(accessToken string, workspaceId string, environmentName string, secretsPath string, includeImports bool) ([]models.SingleEnvironmentVariable, error) {
+func GetPlainTextSecretsViaMachineIdentity(accessToken string, workspaceId string, environmentName string, secretsPath string, includeImports bool) (models.PlaintextSecretResult, error) {
 	httpClient := resty.New()
 	httpClient.SetAuthToken(accessToken).
 		SetHeader("Accept", "application/json")
@@ -170,12 +170,12 @@ func GetPlainTextSecretsViaMachineIdentity(accessToken string, workspaceId strin
 
 	rawSecrets, err := api.CallGetRawSecretsV3(httpClient, api.GetRawSecretsV3Request{WorkspaceId: workspaceId, SecretPath: secretsPath, Environment: environmentName})
 	if err != nil {
-		return nil, err
+		return models.PlaintextSecretResult{}, err
 	}
 
 	plainTextSecrets := []models.SingleEnvironmentVariable{}
 	if err != nil {
-		return nil, fmt.Errorf("unable to decrypt your secrets [err=%v]", err)
+		return models.PlaintextSecretResult{}, fmt.Errorf("unable to decrypt your secrets [err=%v]", err)
 	}
 
 	for _, secret := range rawSecrets.Secrets {
@@ -189,7 +189,10 @@ func GetPlainTextSecretsViaMachineIdentity(accessToken string, workspaceId strin
 	// 	}
 	// }
 
-	return plainTextSecrets, nil
+	return models.PlaintextSecretResult{
+		Secrets: plainTextSecrets,
+		Hash:    rawSecrets.ETag,
+	}, nil
 }
 
 func InjectImportedSecret(plainTextWorkspaceKey []byte, secrets []models.SingleEnvironmentVariable, importedSecrets []api.ImportedSecretV3) ([]models.SingleEnvironmentVariable, error) {