feat(run): support fetching secrets from a non-default organization

packages/cmd/init.go

@@ -4,7 +4,6 @@ Copyright (c) 2023 Infisical Inc.
 package cmd
 
 import (
-	"encoding/json"
 	"fmt"
 
 	"github.com/Infisical/infisical-merge/packages/api"
@@ -135,7 +134,7 @@ var initCmd = &cobra.Command{
 			util.HandleError(err)
 		}
 
-		err = writeWorkspaceFile(filteredWorkspaces[index])
+		err = writeWorkspaceFile(filteredWorkspaces[index], selectedOrgID)
 		if err != nil {
 			util.HandleError(err)
 		}
@@ -211,17 +210,13 @@ func pickOrganization(httpClient *resty.Client, label string, username string) (
 	return subItems[subIndex].ID, &selectedSubOrgName, nil
 }
 
-func writeWorkspaceFile(selectedWorkspace models.Workspace) error {
+func writeWorkspaceFile(selectedWorkspace models.Workspace, organizationId string) error {
 	workspaceFileToSave := models.WorkspaceConfigFile{
-		WorkspaceId: selectedWorkspace.ID,
+		WorkspaceId:    selectedWorkspace.ID,
+		OrganizationId: organizationId,
 	}
 
-	marshalledWorkspaceFile, err := json.MarshalIndent(workspaceFileToSave, "", "    ")
-	if err != nil {
-		return err
-	}
-
-	err = util.WriteToFile(util.INFISICAL_WORKSPACE_CONFIG_FILE_NAME, marshalledWorkspaceFile, 0600)
+	err := util.WriteWorkspaceConfigToPath(workspaceFileToSave, util.INFISICAL_WORKSPACE_CONFIG_FILE_NAME)
 	if err != nil {
 		return err
 	}

packages/cmd/run.go

@@ -82,6 +82,11 @@ var runCmd = &cobra.Command{
 			util.HandleError(err, "Unable to parse flag")
 		}
 
+		organizationId, err := cmd.Flags().GetString("organization-id")
+		if err != nil {
+			util.HandleError(err, "Unable to parse flag")
+		}
+
 		command, err := cmd.Flags().GetString("command")
 		if err != nil {
 			util.HandleError(err, "Unable to parse flag")
@@ -146,15 +151,15 @@ var runCmd = &cobra.Command{
 			ExpandSecretReferences: shouldExpandSecrets,
 		}
 
-		injectableEnvironment, err := fetchAndFormatSecretsForShell(request, projectConfigDir, secretOverriding, token)
+		injectableEnvironment, err := fetchAndFormatSecretsForShell(request, projectConfigDir, secretOverriding, token, organizationId)
 		if err != nil {
 			util.HandleError(err, "Could not fetch secrets", "If you are using a service token to fetch secrets, please ensure it is valid")
 		}
 
 		log.Debug().Msgf("injecting the following environment variables into shell: %v", injectableEnvironment.Variables)
 
 		if watchMode {
-			executeCommandWithWatchMode(command, args, watchModeInterval, request, projectConfigDir, secretOverriding, token)
+			executeCommandWithWatchMode(command, args, watchModeInterval, request, projectConfigDir, secretOverriding, token, organizationId)
 		} else {
 			if cmd.Flags().Changed("command") {
 				command := cmd.Flag("command").Value.String()
@@ -211,6 +216,7 @@ func init() {
 	RootCmd.AddCommand(runCmd)
 	runCmd.Flags().String("token", "", "fetch secrets using service token or machine identity access token")
 	runCmd.Flags().String("projectId", "", "manually set the project ID to fetch secrets from when using machine identity based auth")
+	runCmd.Flags().String("organization-id", "", "manually set the organization ID to fetch secrets from")
 	runCmd.Flags().StringP("env", "e", "dev", "set the environment (dev, prod, etc.) from which your secrets should be pulled from")
 	runCmd.Flags().Bool("expand", true, "parse shell parameter expansions in your secrets")
 	runCmd.Flags().Bool("include-imports", true, "import linked secrets ")
@@ -307,7 +313,7 @@ func waitForExitCommand(cmd *exec.Cmd) (int, error) {
 	return waitStatus.ExitStatus(), nil
 }
 
-func executeCommandWithWatchMode(commandFlag string, args []string, watchModeInterval int, request models.GetMultiPathSecretsParameters, projectConfigDir string, secretOverriding bool, token *models.TokenDetails) {
+func executeCommandWithWatchMode(commandFlag string, args []string, watchModeInterval int, request models.GetMultiPathSecretsParameters, projectConfigDir string, secretOverriding bool, token *models.TokenDetails, organizationId string) {
 
 	var cmd *exec.Cmd
 	var err error
@@ -423,7 +429,7 @@ func executeCommandWithWatchMode(commandFlag string, args []string, watchModeInt
 			watchMutex.Lock()
 			defer watchMutex.Unlock()
 
-			newEnvironmentVariables, err := fetchAndFormatSecretsForShell(request, projectConfigDir, secretOverriding, token)
+			newEnvironmentVariables, err := fetchAndFormatSecretsForShell(request, projectConfigDir, secretOverriding, token, organizationId)
 			if err != nil {
 				log.Error().Err(err).Msg("[HOT RELOAD] Failed to fetch secrets")
 				return
@@ -439,12 +445,13 @@ func executeCommandWithWatchMode(commandFlag string, args []string, watchModeInt
 	}
 }
 
-func fetchSecrets(request models.GetMultiPathSecretsParameters, projectConfigDir string, secretOverriding bool, token *models.TokenDetails) ([]models.SingleEnvironmentVariable, error) {
+func fetchSecrets(request models.GetMultiPathSecretsParameters, projectConfigDir string, secretOverriding bool, token *models.TokenDetails, organizationId string) ([]models.SingleEnvironmentVariable, error) {
 	var allSecrets []models.SingleEnvironmentVariable
 
 	for _, path := range request.SecretsPaths {
 		params := models.GetAllSecretsParameters{
 			Environment:              request.Environment,
+			OrganizationId:           organizationId,
 			WorkspaceId:              request.WorkspaceId,
 			TagSlugs:                 request.TagSlugs,
 			SecretsPath:              path,
@@ -503,8 +510,8 @@ func formatSecretsForShell(secrets []models.SingleEnvironmentVariable) models.In
 	}
 }
 
-func fetchAndFormatSecretsForShell(request models.GetMultiPathSecretsParameters, projectConfigDir string, secretOverriding bool, token *models.TokenDetails) (models.InjectableEnvironmentResult, error) {
-	secrets, err := fetchSecrets(request, projectConfigDir, secretOverriding, token)
+func fetchAndFormatSecretsForShell(request models.GetMultiPathSecretsParameters, projectConfigDir string, secretOverriding bool, token *models.TokenDetails, organizationId string) (models.InjectableEnvironmentResult, error) {
+	secrets, err := fetchSecrets(request, projectConfigDir, secretOverriding, token, organizationId)
 	if err != nil {
 		return models.InjectableEnvironmentResult{}, err
 	}

packages/models/cli.go

@@ -110,6 +110,7 @@ type Workspace struct {
 
 type WorkspaceConfigFile struct {
 	WorkspaceId                   string            `json:"workspaceId"`
+	OrganizationId                string            `json:"organizationId,omitempty"`
 	DefaultEnvironment            string            `json:"defaultEnvironment"`
 	GitBranchToEnvironmentMapping map[string]string `json:"gitBranchToEnvironmentMapping"`
 	Domain                        string            `json:"domain,omitempty"`
@@ -136,6 +137,7 @@ type GetAllSecretsParameters struct {
 	EnvironmentPassedViaFlag bool
 	InfisicalToken           string
 	UniversalAuthAccessToken string
+	OrganizationId           string
 	TagSlugs                 string
 	WorkspaceId              string
 	SecretsPath              string

packages/util/config.go

@@ -197,6 +197,20 @@ func GetWorkspaceConfigByPath(path string) (workspaceConfig models.WorkspaceConf
 	return workspaceConfigFile, nil
 }
 
+func WriteWorkspaceConfigToPath(config models.WorkspaceConfigFile, path string) error {
+	marshalledWorkspaceFile, err := json.MarshalIndent(config, "", "    ")
+	if err != nil {
+		return err
+	}
+
+	err = WriteToFile(path, marshalledWorkspaceFile, 0600)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // Get the infisical config file and if it doesn't exist, return empty config model, otherwise raise error
 func GetConfigFile() (models.ConfigFile, error) {
 	fullConfigFilePath, _, err := GetFullConfigFilePath()

packages/util/secrets.go

@@ -8,13 +8,15 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"path/filepath"
 	"strings"
 	"unicode"
 
 	"github.com/Infisical/infisical-merge/packages/api"
 	"github.com/Infisical/infisical-merge/packages/crypto"
 	"github.com/Infisical/infisical-merge/packages/models"
 	"github.com/go-resty/resty/v2"
+	jwt "github.com/golang-jwt/jwt/v5"
 	"github.com/rs/zerolog/log"
 	"github.com/zalando/go-keyring"
 	"gopkg.in/yaml.v3"
@@ -280,6 +282,8 @@ func GetAllEnvironmentVariables(params models.GetAllSecretsParameters, projectCo
 	var secretsToReturn []models.SingleEnvironmentVariable
 	// var serviceTokenDetails api.GetServiceTokenDetailsResponse
 	var errorToReturn error
+	var workspaceConfigFile models.WorkspaceConfigFile
+	var workspaceConfigFilePath string
 
 	if params.InfisicalToken == "" && params.UniversalAuthAccessToken == "" {
 		if params.WorkspaceId == "" {
@@ -313,28 +317,50 @@ func GetAllEnvironmentVariables(params models.GetAllSecretsParameters, projectCo
 		}
 
 		if params.WorkspaceId == "" {
-			var infisicalDotJson models.WorkspaceConfigFile
-
 			if projectConfigFilePath == "" {
+				workspaceConfigFilePath, err = FindWorkspaceConfigFile()
+				if err != nil {
+					PrintErrorMessageAndExit("Please either run infisical init to connect to a project or pass in project id with --projectId flag")
+				}
 				projectConfig, err := GetWorkSpaceFromFile()
 				if err != nil {
 					PrintErrorMessageAndExit("Please either run infisical init to connect to a project or pass in project id with --projectId flag")
 				}
 
-				infisicalDotJson = projectConfig
+				workspaceConfigFile = projectConfig
 			} else {
+				workspaceConfigFilePath = filepath.Join(projectConfigFilePath, INFISICAL_WORKSPACE_CONFIG_FILE_NAME)
 				projectConfig, err := GetWorkSpaceFromFilePath(projectConfigFilePath)
 				if err != nil {
 					return nil, err
 				}
 
-				infisicalDotJson = projectConfig
+				workspaceConfigFile = projectConfig
+			}
+			params.WorkspaceId = workspaceConfigFile.WorkspaceId
+		} else if projectConfigFilePath == "" {
+			workspaceConfigFilePath, err = FindWorkspaceConfigFile()
+			if err != nil {
+				workspaceConfigFilePath = ""
+			}
+			projectConfig, err := GetWorkSpaceFromFile()
+			if err == nil {
+				workspaceConfigFile = projectConfig
+			}
+		} else {
+			workspaceConfigFilePath = filepath.Join(projectConfigFilePath, INFISICAL_WORKSPACE_CONFIG_FILE_NAME)
+			projectConfig, err := GetWorkSpaceFromFilePath(projectConfigFilePath)
+			if err == nil {
+				workspaceConfigFile = projectConfig
 			}
-			params.WorkspaceId = infisicalDotJson.WorkspaceId
 		}
 
-		res, err := GetPlainTextSecretsV4(loggedInUserDetails.UserCredentials.JTWToken, params.WorkspaceId,
-			params.Environment, params.SecretsPath, params.IncludeImport, params.Recursive, params.TagSlugs, true, params.IncludePersonalOverrides)
+		resolvedToken, err := resolveOrgScopedToken(loggedInUserDetails, params.OrganizationId, workspaceConfigFile.OrganizationId)
+		if err != nil {
+			return nil, err
+		}
+
+		res, err := fetchSecretsWithOrgDiscovery(loggedInUserDetails, resolvedToken, params, workspaceConfigFile, workspaceConfigFilePath)
 		log.Debug().Msgf("GetAllEnvironmentVariables: Trying to fetch secrets JTW token [err=%s]", err)
 
 		if err == nil {
@@ -381,6 +407,129 @@ func GetAllEnvironmentVariables(params models.GetAllSecretsParameters, projectCo
 	return secretsToReturn, errorToReturn
 }
 
+type tokenOrganizationClaims struct {
+	OrganizationId string `json:"organizationId"`
+	jwt.RegisteredClaims
+}
+
+func getTokenOrganizationId(token string) (string, error) {
+	var claims tokenOrganizationClaims
+	parser := jwt.NewParser()
+	if _, _, err := parser.ParseUnverified(token, &claims); err != nil {
+		return "", err
+	}
+
+	return claims.OrganizationId, nil
+}
+
+func resolveOrgScopedToken(loggedInUserDetails LoggedInUserDetails, flagOrganizationId string, configOrganizationId string) (string, error) {
+	targetOrgId := flagOrganizationId
+	if targetOrgId == "" {
+		targetOrgId = os.Getenv("INFISICAL_ORGANIZATION_ID")
+	}
+	if targetOrgId == "" {
+		targetOrgId = configOrganizationId
+	}
+	if targetOrgId == "" {
+		return loggedInUserDetails.UserCredentials.JTWToken, nil
+	}
+
+	currentTokenOrgId, err := getTokenOrganizationId(loggedInUserDetails.UserCredentials.JTWToken)
+	if err != nil {
+		return "", fmt.Errorf("unable to determine the organization scope for the current login token [err=%v]", err)
+	}
+	if currentTokenOrgId == targetOrgId {
+		return loggedInUserDetails.UserCredentials.JTWToken, nil
+	}
+
+	httpClient, err := GetRestyClientWithCustomHeaders()
+	if err != nil {
+		return "", fmt.Errorf("unable to get resty client with custom headers [err=%v]", err)
+	}
+	httpClient.SetAuthToken(loggedInUserDetails.UserCredentials.JTWToken)
+
+	selectOrganizationResponse, err := api.CallSelectOrganization(httpClient, api.SelectOrganizationRequest{OrganizationId: targetOrgId})
+	if err != nil {
+		return "", fmt.Errorf("unable to scope the current session to organization %s; ensure your account can access that organization or log in again [err=%v]", targetOrgId, err)
+	}
+
+	return selectOrganizationResponse.Token, nil
+}
+
+func isOrganizationScopeError(err error) bool {
+	var apiErr *api.APIError
+	if !errors.As(err, &apiErr) || apiErr.StatusCode != 403 {
+		return false
+	}
+
+	// Only the org-scoping rejection should trigger discovery; a generic
+	// permission 403 must not enumerate every organization.
+	return strings.Contains(apiErr.ErrorMessage, "does not belong to your selected organization")
+}
+
+func fetchSecretsWithOrgDiscovery(loggedInUserDetails LoggedInUserDetails, resolvedToken string, params models.GetAllSecretsParameters, workspaceConfigFile models.WorkspaceConfigFile, workspaceConfigFilePath string) (models.PlaintextSecretResult, error) {
+	res, err := GetPlainTextSecretsV4(resolvedToken, params.WorkspaceId,
+		params.Environment, params.SecretsPath, params.IncludeImport, params.Recursive, params.TagSlugs, true, params.IncludePersonalOverrides)
+	if err == nil {
+		return res, nil
+	}
+	if !isOrganizationScopeError(err) {
+		return models.PlaintextSecretResult{}, err
+	}
+
+	originalError := err
+	httpClient, err := GetRestyClientWithCustomHeaders()
+	if err != nil {
+		return models.PlaintextSecretResult{}, originalError
+	}
+	httpClient.SetAuthToken(loggedInUserDetails.UserCredentials.JTWToken)
+
+	orgsResponse, err := api.CallGetAllOrganizations(httpClient)
+	if err != nil {
+		return models.PlaintextSecretResult{}, originalError
+	}
+
+	currentTokenOrgId, err := getTokenOrganizationId(loggedInUserDetails.UserCredentials.JTWToken)
+	if err != nil {
+		return models.PlaintextSecretResult{}, originalError
+	}
+
+	for _, organization := range orgsResponse.Organizations {
+		if organization.ID == currentTokenOrgId {
+			continue
+		}
+
+		httpClient.SetAuthToken(loggedInUserDetails.UserCredentials.JTWToken)
+		selectOrganizationResponse, err := api.CallSelectOrganization(httpClient, api.SelectOrganizationRequest{OrganizationId: organization.ID})
+		if err != nil {
+			continue
+		}
+
+		res, err := GetPlainTextSecretsV4(selectOrganizationResponse.Token, params.WorkspaceId,
+			params.Environment, params.SecretsPath, params.IncludeImport, params.Recursive, params.TagSlugs, true, params.IncludePersonalOverrides)
+		if err != nil {
+			continue
+		}
+
+		persistDiscoveredOrganizationId(workspaceConfigFile, workspaceConfigFilePath, organization.ID)
+		return res, nil
+	}
+
+	return models.PlaintextSecretResult{}, originalError
+}
+
+func persistDiscoveredOrganizationId(workspaceConfigFile models.WorkspaceConfigFile, workspaceConfigFilePath string, organizationId string) {
+	if workspaceConfigFilePath == "" {
+		return
+	}
+
+	workspaceConfigFile.OrganizationId = organizationId
+	err := WriteWorkspaceConfigToPath(workspaceConfigFile, workspaceConfigFilePath)
+	if err != nil {
+		log.Debug().Err(err).Str("path", workspaceConfigFilePath).Msg("Failed to persist discovered organization ID to workspace config")
+	}
+}
+
 func GetBackupEncryptionKey() ([]byte, error) {
 	encryptionKey, err := GetValueInKeyring(INFISICAL_BACKUP_SECRET_ENCRYPTION_KEY)
 	if err != nil {