From 127e1902cca9159a0d19ebfd944c5ee163041fd0 Mon Sep 17 00:00:00 2001 From: "Baruch Odem (Rothkoff)" Date: Thu, 29 Jun 2023 14:49:16 +0300 Subject: [PATCH] feat: option to load arguments from config file (#107) - bind flags from viper to cobra command - add tests - test for Execute - support multiple config formats - update readme with help message - document configuration env and file Closes #21 **Proposed Changes** - feat: read arguments from Environment Variables - feat: read arguments from JSON/YAML file - docs: auto-update the README when the `help` message was changed - docs: how to use arguments from Environment Variables or config file **Limitations** - You still can run only one plugin at a time. - You have to add the plugin command to the CLI command. - Positional arguments are not yet supported. --------- Co-authored-by: Jossef Harush Kadouri --- .github/workflows/pr-validation.yml | 10 +- README.md | 73 ++- cmd/main.go | 26 +- go.mod | 7 +- lib/flags.go | 101 ++++ lib/flags_test.go | 735 ++++++++++++++++++++++++++++ 6 files changed, 938 insertions(+), 14 deletions(-) create mode 100644 lib/flags.go create mode 100644 lib/flags_test.go diff --git a/.github/workflows/pr-validation.yml b/.github/workflows/pr-validation.yml index 4decdc74..6f998c8a 100644 --- a/.github/workflows/pr-validation.yml +++ b/.github/workflows/pr-validation.yml @@ -21,17 +21,17 @@ jobs: with: go-version: "^1.20" + - name: go mod tidy + run: | + go mod tidy + git diff --exit-code + - name: Go Linter run: docker run --rm -v $(pwd):/app -w /app golangci/golangci-lint:v1.52.0 golangci-lint run -v -E gofmt --timeout=5m - name: Go Test run: go test -v ./... - - name: go mod tidy - run: | - go mod tidy - git diff --exit-code - build: runs-on: ubuntu-latest steps: diff --git a/README.md b/README.md index ed6e8fff..a04afb82 100644 --- a/README.md +++ b/README.md @@ -36,12 +36,75 @@ docker run -v path/to/my/repo:/repo checkmarx/2ms git /repo ## Getting started -### Command line arguments (wip, see [#20](https://github.com/Checkmarx/2ms/discussions/20)) +``` +2ms Secrets Detection: A tool to detect secrets in public websites and communication services. + +Usage: + 2ms [command] + +Commands + confluence Scan Confluence server + discord Scan Discord server + filesystem Scan local folder + git Scan Git repository + paligo Scan Paligo instance + slack Scan Slack team + +Additional Commands: + completion Generate the autocompletion script for the specified shell + help Help about any command + +Flags: + --config string YAML config file path + -h, --help help for 2ms + --log-level string log level (trace, debug, info, warn, error, fatal) (default "info") + --regex stringArray custom regexes to apply to the scan, must be valid Go regex + --report-path strings path to generate report files. The output format will be determined by the file extension (.json, .yaml, .sarif) + --stdout-format string stdout output format, available formats are: json, yaml, sarif (default "yaml") + --tags strings select rules to be applied (default [all]) + -v, --version version for 2ms + +Use "2ms [command] --help" for more information about a command. +``` + +| :warning: Using configuration env or file | +| :--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Please note that even using configuration file or environment variables, you still need to specify the subcommand name in the CLI arguments. Also, positional arguments are not yet supported. | + +### Environment Variables + +To use a flag as an environment variable, see the following rules: + +- Replace `-` with `_` +- Start with `2MS_` +- Prefer uppercase +- Append the subcommand name(s) (if any) with `_` -- `--confluence` The URL of the Confluence instance to scan. -- `--confluence-spaces` A comma-separated list of Confluence spaces to scan. -- `--confluence-username` confluence username or email -- `--confluence-token` confluence token +Examples: + +- `--log-level` -> `2MS_LOG_LEVEL` +- `paligo instance` -> `2MS_PALIGO_INSTANCE` + +### Configuration File + +You can use `--config` flag to specify a configuration file. The configuration file is a YAML/JSON file with the following structure: + +```yaml +# global flags that will be applied to all commands +log-level: info +report-path: + - ./report.yaml + - ./report.json + - ./report.sarif + +# the subcommand will be selected from the CLI arguments +# the flags below will be applied to the selected subcommand +paligo: + instance: your-instance + username: your-username + # you can combine config file and Environment Variables + # token: your-token +``` --- diff --git a/cmd/main.go b/cmd/main.go index 96342b6e..a01d98e1 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -7,6 +7,7 @@ import ( "strings" "github.com/checkmarx/2ms/config" + "github.com/checkmarx/2ms/lib" "sync" "time" @@ -18,6 +19,7 @@ import ( "github.com/rs/zerolog" "github.com/rs/zerolog/log" "github.com/spf13/cobra" + "github.com/spf13/viper" ) var Version = "0.0.0" @@ -27,6 +29,7 @@ const ( jsonFormat = "json" yamlFormat = "yaml" sarifFormat = "sarif" + configFileFlag = "config" logLevelFlagName = "log-level" reportPathFlagName = "report-path" @@ -52,6 +55,11 @@ var rootCmd = &cobra.Command{ Version: Version, } +const envPrefix = "2MS" + +var configFilePath string +var vConfig = viper.New() + var allPlugins = []plugins.IPlugin{ &plugins.ConfluencePlugin{}, &plugins.DiscordPlugin{}, @@ -70,7 +78,16 @@ var channels = plugins.Channels{ var report = reporting.Init() var secretsChan = make(chan reporting.Secret) -func initLog() { +func initialize() { + zerolog.SetGlobalLevel(zerolog.InfoLevel) + + configFilePath, err := rootCmd.Flags().GetString(configFileFlag) + if err != nil { + cobra.CheckErr(err) + } + cobra.CheckErr(lib.LoadConfig(vConfig, configFilePath)) + cobra.CheckErr(lib.BindFlags(rootCmd, vConfig, envPrefix)) + switch strings.ToLower(logLevelVar) { case "trace": zerolog.SetGlobalLevel(zerolog.TraceLevel) @@ -90,7 +107,12 @@ func initLog() { } func Execute() { - cobra.OnInitialize(initLog) + vConfig.SetEnvPrefix(envPrefix) + vConfig.AutomaticEnv() + + cobra.OnInitialize(initialize) + rootCmd.PersistentFlags().StringVar(&configFilePath, configFileFlag, "", "config file path") + cobra.CheckErr(rootCmd.MarkPersistentFlagFilename(configFileFlag, "yaml", "yml", "json")) rootCmd.PersistentFlags().StringVar(&logLevelVar, logLevelFlagName, "info", "log level (trace, debug, info, warn, error, fatal)") rootCmd.PersistentFlags().StringSliceVar(&reportPathVar, reportPathFlagName, []string{}, "path to generate report files. The output format will be determined by the file extension (.json, .yaml, .sarif)") rootCmd.PersistentFlags().StringVar(&stdoutFormatVar, stdoutFormatFlagName, "yaml", "stdout output format, available formats are: json, yaml, sarif") diff --git a/go.mod b/go.mod index 53783422..7846b375 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,9 @@ require ( github.com/rs/zerolog v1.29.0 github.com/slack-go/slack v0.12.2 github.com/spf13/cobra v1.6.1 + github.com/spf13/pflag v1.0.5 + github.com/spf13/viper v1.15.0 + github.com/stretchr/testify v1.8.1 github.com/zricethezav/gitleaks/v8 v8.16.1 golang.org/x/time v0.1.0 gopkg.in/yaml.v2 v2.4.0 @@ -16,6 +19,7 @@ require ( require ( github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect github.com/charmbracelet/lipgloss v0.7.1 // indirect + github.com/davecgh/go-spew v1.1.1 // indirect github.com/fatih/semgroup v1.2.0 // indirect github.com/fsnotify/fsnotify v1.6.0 // indirect github.com/gorilla/websocket v1.5.0 // indirect @@ -33,12 +37,11 @@ require ( github.com/muesli/termenv v0.15.1 // indirect github.com/pelletier/go-toml/v2 v2.0.7 // indirect github.com/petar-dambovaliev/aho-corasick v0.0.0-20211021192214-5ab2d9280aa9 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect github.com/rivo/uniseg v0.4.4 // indirect github.com/spf13/afero v1.9.5 // indirect github.com/spf13/cast v1.5.0 // indirect github.com/spf13/jwalterweatherman v1.1.0 // indirect - github.com/spf13/pflag v1.0.5 // indirect - github.com/spf13/viper v1.15.0 // indirect github.com/subosito/gotenv v1.4.2 // indirect golang.org/x/crypto v0.0.0-20220722155217-630584e8d5aa // indirect golang.org/x/sync v0.1.0 // indirect diff --git a/lib/flags.go b/lib/flags.go new file mode 100644 index 00000000..b9a3fbea --- /dev/null +++ b/lib/flags.go @@ -0,0 +1,101 @@ +package lib + +import ( + "fmt" + "path/filepath" + "strings" + + "github.com/rs/zerolog/log" + "github.com/spf13/cobra" + "github.com/spf13/pflag" + "github.com/spf13/viper" +) + +func LoadConfig(v *viper.Viper, configFilePath string) error { + if configFilePath == "" { + return nil + } + + configType := strings.TrimPrefix(filepath.Ext(configFilePath), ".") + + v.SetConfigType(configType) + v.SetConfigFile(configFilePath) + return v.ReadInConfig() +} + +// TODO: can be a package + +// BindFlags fill flags values with config file or environment variables data +func BindFlags(cmd *cobra.Command, v *viper.Viper, envPrefix string) error { + commandHierarchy := getCommandHierarchy(cmd) + + bindFlag := func(f *pflag.Flag) { + fullFlagName := fmt.Sprintf("%s%s", commandHierarchy, f.Name) + bindEnvVarIntoViper(v, fullFlagName, envPrefix) + + if f.Changed { + return + } + + if v.IsSet(fullFlagName) { + val := v.Get(fullFlagName) + applyViperFlagToCommand(f, val, cmd) + } + } + cmd.PersistentFlags().VisitAll(bindFlag) + cmd.Flags().VisitAll(bindFlag) + + for _, subCmd := range cmd.Commands() { + if err := BindFlags(subCmd, v, envPrefix); err != nil { + return err + } + } + + return nil +} + +func bindEnvVarIntoViper(v *viper.Viper, fullFlagName, envPrefix string) { + envVarSuffix := strings.ToUpper(strings.ReplaceAll(strings.ReplaceAll(fullFlagName, "-", "_"), ".", "_")) + envVarName := fmt.Sprintf("%s_%s", envPrefix, envVarSuffix) + + if err := v.BindEnv(fullFlagName, envVarName, strings.ToLower(envVarName)); err != nil { + log.Err(err).Msg("Failed to bind Viper flags") + } +} + +func applyViperFlagToCommand(flag *pflag.Flag, val interface{}, cmd *cobra.Command) { + switch t := val.(type) { + case []interface{}: + var paramSlice []string + for _, param := range t { + paramSlice = append(paramSlice, param.(string)) + } + valStr := strings.Join(paramSlice, ",") + if err := flag.Value.Set(valStr); err != nil { + log.Err(err).Msg("Failed to set Viper flags") + } + default: + newVal := fmt.Sprintf("%v", val) + if err := flag.Value.Set(newVal); err != nil { + log.Err(err).Msg("Failed to set Viper flags") + } + } + flag.Changed = true +} + +func getCommandHierarchy(cmd *cobra.Command) string { + names := []string{} + if !cmd.HasParent() { + return "" + } + + for parent := cmd; parent.HasParent() && parent.Name() != ""; parent = parent.Parent() { + names = append([]string{parent.Name()}, names...) + } + + if len(names) == 0 { + return "" + } + + return strings.Join(names, ".") + "." +} diff --git a/lib/flags_test.go b/lib/flags_test.go new file mode 100644 index 00000000..b8b82f44 --- /dev/null +++ b/lib/flags_test.go @@ -0,0 +1,735 @@ +package lib_test + +import ( + "bytes" + "os" + "strings" + "testing" + + "github.com/checkmarx/2ms/lib" + "github.com/spf13/cobra" + "github.com/spf13/viper" + "github.com/stretchr/testify/assert" +) + +const envVarPrefix = "PREFIX" + +func TestBindFlags(t *testing.T) { + t.Run("BindFlags_TestEmptyViper", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + cmd := &cobra.Command{} + v := getViper() + + var ( + testString string + testInt int + testBool bool + testFloat64 float64 + ) + + cmd.PersistentFlags().StringVar(&testString, "test-string", "", "Test string flag") + cmd.PersistentFlags().IntVar(&testInt, "test-int", 0, "Test int flag") + cmd.PersistentFlags().BoolVar(&testBool, "test-bool", false, "Test bool flag") + cmd.PersistentFlags().Float64Var(&testFloat64, "test-float64", 0.0, "Test float64 flag") + + err := lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Empty(t, testString) + assert.Empty(t, testInt) + assert.Empty(t, testBool) + assert.Empty(t, testFloat64) + }) + + t.Run("BindFlags_FromEnvVarsToCobraCommand", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + cmd := &cobra.Command{} + v := getViper() + v.SetEnvPrefix(envVarPrefix) + + var ( + testString string + testInt int + testBool bool + testFloat64 float64 + ) + + cmd.PersistentFlags().StringVar(&testString, "test-string", "", "Test string flag") + cmd.PersistentFlags().IntVar(&testInt, "test-int", 0, "Test int flag") + cmd.PersistentFlags().BoolVar(&testBool, "test-bool", false, "Test bool flag") + cmd.PersistentFlags().Float64Var(&testFloat64, "test-float64", 0.0, "Test float64 flag") + + err := setEnv("PREFIX_TEST_STRING", "test-string-value") + assert.NoError(t, err) + err = setEnv("PREFIX_TEST_INT", "456") + assert.NoError(t, err) + err = setEnv("PREFIX_TEST_BOOL", "true") + assert.NoError(t, err) + err = setEnv("PREFIX_TEST_FLOAT64", "1.23") + assert.NoError(t, err) + + err = lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "test-string-value", testString) + assert.Equal(t, 456, testInt) + assert.Equal(t, true, testBool) + assert.Equal(t, 1.23, testFloat64) + }) + + t.Run("BindFlags_NonPersistentFlags", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + cmd := &cobra.Command{} + v := getViper() + + var ( + testString string + ) + + cmd.Flags().StringVar(&testString, "test-string", "", "Test string flag") + + err := setEnv("PREFIX_TEST_STRING", "test-string-value") + assert.NoError(t, err) + + err = lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "test-string-value", testString) + }) + + t.Run("BindFlags_Subcommand", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + var ( + testString string + testInt int + ) + + subCommand := &cobra.Command{ + Use: "subCommand", + } + subCommand.Flags().StringVar(&testString, "test-string", "", "Test string flag") + subCommand.PersistentFlags().IntVar(&testInt, "test-int", 0, "Test int flag") + + cmd := &cobra.Command{} + cmd.AddCommand(subCommand) + v := getViper() + + err := setEnv("PREFIX_SUBCOMMAND_TEST_STRING", "test-string-value") + assert.NoError(t, err) + err = setEnv("PREFIX_SUBCOMMAND_TEST_INT", "456") + assert.NoError(t, err) + + err = lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "test-string-value", testString) + assert.Equal(t, 456, testInt) + }) + + t.Run("BindFlags_ArrayFlag", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + arr := []string{"test", "array", "flag"} + + cmd := &cobra.Command{} + v := getViper() + + var ( + // testArraySpaces []string + testArrayCommas []string + ) + + // cmd.PersistentFlags().StringSliceVar(&testArraySpaces, "test-array-spaces", []string{}, "Test array flag") + cmd.PersistentFlags().StringSliceVar(&testArrayCommas, "test-array-commas", []string{}, "Test array flag") + + // err := setEnv("PREFIX_TEST_ARRAY_SPACES", strings.Join(arr, " ")) + // assert.NoError(t, err) + err := setEnv("PREFIX_TEST_ARRAY_COMMAS", strings.Join(arr, ",")) + assert.NoError(t, err) + + err = lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + // assert.Equal(t, testArraySpaces, arr) + assert.Equal(t, arr, testArrayCommas) + }) + + t.Run("BindFlags_ReturnsErrorForUnknownConfigurationKeys", func(t *testing.T) { + t.Skip("Not sure if we need this feature.") + assertClearEnv(t) + defer clearEnvVars(t) + + cmd := &cobra.Command{} + v := getViper() + + var ( + testString string + ) + + cmd.PersistentFlags().StringVar(&testString, "test-string", "", "Test string flag") + + v.Set("unknown-key", "unknown-value") + + err := lib.BindFlags(cmd, v, envVarPrefix) + + assert.EqualError(t, err, "unknown configuration key: 'unknown-key'\nShowing help for '' command") + }) + + t.Run("BindFlags_LowerCaseEnvVars", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + cmd := &cobra.Command{} + v := getViper() + + var ( + testString string + ) + + cmd.PersistentFlags().StringVar(&testString, "test-string", "", "Test string flag") + + err := setEnv("prefix_test_string", "test-string-value") + assert.NoError(t, err) + + err = lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "test-string-value", testString) + }) + + t.Run("BindFlags_OneWordFlagName", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + cmd := &cobra.Command{} + v := getViper() + + var ( + testString string + ) + + cmd.Flags().StringVar(&testString, "teststring", "", "Test string flag") + + err := setEnv("prefix_teststring", "test-string-value") + assert.NoError(t, err) + + err = lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "test-string-value", testString) + }) + + t.Run("BindFlags_SameFlagNameDifferentCmd", func(t *testing.T) { + + assertClearEnv(t) + defer clearEnvVars(t) + + rootCmd := &cobra.Command{ + Use: "root", + } + cmd1 := &cobra.Command{ + Use: "cmd1", + } + cmd2 := &cobra.Command{ + Use: "cmd2", + } + v := getViper() + + var ( + testStringRoot string + testStringPersistentRoot string + testString1 string + testStringPersistent1 string + testString2 string + testStringPersistent2 string + ) + + rootCmd.Flags().StringVar(&testStringRoot, "test-string", "", "Test string flag") + rootCmd.PersistentFlags().StringVar(&testStringPersistentRoot, "test-string-persistent", "", "Test string flag") + cmd1.Flags().StringVar(&testString1, "test-string", "", "Test string flag") + cmd1.PersistentFlags().StringVar(&testStringPersistent1, "test-string-persistent", "", "Test string flag") + cmd2.Flags().StringVar(&testString2, "test-string", "", "Test string flag") + cmd2.PersistentFlags().StringVar(&testStringPersistent2, "test-string-persistent", "", "Test string flag") + + rootCmd.AddCommand(cmd1) + rootCmd.AddCommand(cmd2) + + err := setEnv("prefix_test_string", "test-string-value") + assert.NoError(t, err) + err = setEnv("prefix_test_string_persistent", "test-string-persistent-value") + assert.NoError(t, err) + err = setEnv("prefix_cmd1_test_string", "test-string-value-cmd1") + assert.NoError(t, err) + err = setEnv("prefix_cmd1_test_string_persistent", "test-string-persistent-value-cmd1") + assert.NoError(t, err) + err = setEnv("prefix_cmd2_test_string", "test-string-value-cmd2") + assert.NoError(t, err) + err = setEnv("prefix_cmd2_test_string_persistent", "test-string-persistent-value-cmd2") + assert.NoError(t, err) + + err = lib.BindFlags(rootCmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "test-string-value", testStringRoot) + assert.Equal(t, "test-string-persistent-value", testStringPersistentRoot) + assert.Equal(t, "test-string-value-cmd1", testString1) + assert.Equal(t, "test-string-persistent-value-cmd1", testStringPersistent1) + assert.Equal(t, "test-string-value-cmd2", testString2) + assert.Equal(t, "test-string-persistent-value-cmd2", testStringPersistent2) + }) + + t.Run("BindFlags_FromYAML_RootCMD", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + yamlConfig := []byte(` +test-string: test-string-value +test-int: 123 +test-bool: true +test-array: + - test + - array + - flag +test-float: 123.456 +`) + + cmd := &cobra.Command{} + v := getViper() + v.SetConfigType("yaml") + assert.NoError(t, v.ReadConfig(bytes.NewBuffer(yamlConfig))) + + var ( + testString string + testInt int + testBool bool + testArray []string + testFloat float64 + ) + + cmd.PersistentFlags().StringVar(&testString, "test-string", "", "Test string flag") + cmd.Flags().IntVar(&testInt, "test-int", 0, "Test int flag") + cmd.PersistentFlags().BoolVar(&testBool, "test-bool", false, "Test bool flag") + cmd.Flags().StringSliceVar(&testArray, "test-array", []string{}, "Test array flag") + cmd.PersistentFlags().Float64Var(&testFloat, "test-float", 0, "Test float flag") + + err := lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "test-string-value", testString) + assert.Equal(t, 123, testInt) + assert.Equal(t, true, testBool) + assert.Equal(t, []string{"test", "array", "flag"}, testArray) + assert.Equal(t, 123.456, testFloat) + }) + + t.Run("BindFlags_FromYAML_SubCMD", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + yamlConfig := []byte(` +global-string: global-string-value +subCommand: + test-string: test-string-value + test-int: 123 + test-bool: true +`) + + cmd := &cobra.Command{} + v := getViper() + v.SetConfigType("yaml") + assert.NoError(t, v.ReadConfig(bytes.NewBuffer(yamlConfig))) + + var ( + globalString string + testString string + testInt int + testBool bool + ) + + cmd.PersistentFlags().StringVar(&globalString, "global-string", "", "Global string flag") + subCmd := &cobra.Command{ + Use: "subCommand", + } + cmd.AddCommand(subCmd) + subCmd.PersistentFlags().StringVar(&testString, "test-string", "", "Test string flag") + subCmd.Flags().IntVar(&testInt, "test-int", 0, "Test int flag") + subCmd.PersistentFlags().BoolVar(&testBool, "test-bool", false, "Test bool flag") + + err := lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "global-string-value", globalString) + assert.Equal(t, "test-string-value", testString) + assert.Equal(t, 123, testInt) + assert.Equal(t, true, testBool) + }) + + t.Run("BindFlags_FromYAML_SubCMD_WithEnvVars", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + yamlConfig := []byte(` +global-string: global-string-value +subCommand: + test-string: test-string-value + test-int: 123 + test-bool: true +`) + cmd := &cobra.Command{} + v := getViper() + v.SetConfigType("yaml") + assert.NoError(t, v.ReadConfig(bytes.NewBuffer(yamlConfig))) + + var ( + globalString string + testString string + testInt int + testBool bool + ) + + cmd.PersistentFlags().StringVar(&globalString, "global-string", "", "Global string flag") + subCmd := &cobra.Command{ + Use: "subCommand", + } + cmd.AddCommand(subCmd) + subCmd.PersistentFlags().StringVar(&testString, "test-string", "", "Test string flag") + subCmd.Flags().IntVar(&testInt, "test-int", 0, "Test int flag") + subCmd.PersistentFlags().BoolVar(&testBool, "test-bool", false, "Test bool flag") + + err := setEnv("PREFIX_GLOBAL_STRING", "global-string-value-from-env") + assert.NoError(t, err) + err = setEnv("PREFIX_SUBCOMMAND_TEST_STRING", "test-string-value-from-env") + assert.NoError(t, err) + + err = lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "global-string-value-from-env", globalString) + assert.Equal(t, "test-string-value-from-env", testString) + assert.Equal(t, 123, testInt) + assert.Equal(t, true, testBool) + }) + + t.Run("BindFlags_FromYAML_SubSubCmd", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + yamlConfig := []byte(` +global-string: global-string-value +subCommand: + first-string: string-from-sub-command + subSubCommand: + second-string: string from sub-sub command +`) + cmd := &cobra.Command{} + v := getViper() + v.SetConfigType("yaml") + assert.NoError(t, v.ReadConfig(bytes.NewBuffer(yamlConfig))) + + var ( + globalString string + firstString string + secondString string + ) + + subSubCmd := &cobra.Command{ + Use: "subSubCommand", + } + subCmd := &cobra.Command{ + Use: "subCommand", + } + subCmd.AddCommand(subSubCmd) + cmd.AddCommand(subCmd) + cmd.PersistentFlags().StringVar(&globalString, "global-string", "", "Global string flag") + subCmd.PersistentFlags().StringVar(&firstString, "first-string", "", "Test string flag") + subSubCmd.Flags().StringVar(&secondString, "second-string", "", "Test string flag") + + err := lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "global-string-value", globalString) + assert.Equal(t, "string-from-sub-command", firstString) + assert.Equal(t, "string from sub-sub command", secondString) + }) + + t.Run("BindFlags_FromYAML_SameFlagName_Root", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + yamlConfig := []byte(` +test-string: global-string-value +subCommand: + dummy-string: string-from-sub-command +`) + + cmd := &cobra.Command{} + v := getViper() + v.SetConfigType("yaml") + assert.NoError(t, v.ReadConfig(bytes.NewBuffer(yamlConfig))) + + var ( + testStringRoot string + testStringSub string + ) + + subCmd := &cobra.Command{ + Use: "subCommand", + } + cmd.AddCommand(subCmd) + + cmd.PersistentFlags().StringVar(&testStringRoot, "test-string", "", "Test string flag") + subCmd.PersistentFlags().StringVar(&testStringSub, "test-string", "", "Test string flag") + + err := lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "global-string-value", testStringRoot) + assert.Equal(t, "", testStringSub) + }) + + t.Run("BindFlags_FromYAML_SameFlagName_SubCmd", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + yamlConfig := []byte(` +test-string: global-string-value +subCommand: + test-string: string-from-sub-command +`) + + cmd := &cobra.Command{} + v := getViper() + v.SetConfigType("yaml") + assert.NoError(t, v.ReadConfig(bytes.NewBuffer(yamlConfig))) + + var ( + testStringRoot string + testStringSub string + ) + + subCmd := &cobra.Command{ + Use: "subCommand", + } + + cmd.PersistentFlags().StringVar(&testStringRoot, "test-string", "", "Test string flag") + subCmd.PersistentFlags().StringVar(&testStringSub, "test-string", "", "Test string flag") + + cmd.AddCommand(subCmd) + + err := lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "global-string-value", testStringRoot) + assert.Equal(t, "string-from-sub-command", testStringSub) + }) + + t.Run("BindFlags_FromJSON", func(t *testing.T) { + assertClearEnv(t) + defer clearEnvVars(t) + + jsonConfig := []byte(` + { + "global-string": "global-string-value", + "subCommand": { + "test-string": "string-from-sub-command" + } + }`) + + cmd := &cobra.Command{} + v := getViper() + v.SetConfigType("json") + assert.NoError(t, v.ReadConfig(bytes.NewBuffer(jsonConfig))) + + subCmd := &cobra.Command{ + Use: "subCommand", + } + cmd.AddCommand(subCmd) + + globalString := cmd.PersistentFlags().String("global-string", "", "Global string flag") + testString := subCmd.PersistentFlags().String("test-string", "", "Test string flag") + + err := lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + + assert.Equal(t, "global-string-value", *globalString) + assert.Equal(t, "string-from-sub-command", *testString) + }) +} + +func TestEndToEndWithExecute(t *testing.T) { + configFlagName := "config" + + testCases := []struct { + name string + args []string + envVars map[string]string + config []byte + configFormat string + }{ + { + name: "from env vars", + args: []string{"subcommand"}, + envVars: map[string]string{"TEST_STRING": "env-value", "TEST_INT": "123", "SUBCOMMAND_TEST_BOOL": "true"}, + }, + { + name: "from argument", + args: []string{"subcommand", "--test-string", "argument-value", "--test-int", "123", "--test-bool", "true"}, + }, + { + name: "from config", + args: []string{"subcommand"}, + config: []byte(` +test-string: config-value +test-int: 123 +subcommand: + test-bool: true +`), + configFormat: "yaml", + }, + { + name: "from argument and env vars", + args: []string{"subcommand", "--test-string", "argument-value"}, + envVars: map[string]string{ + "TEST_INT": "123", + "SUBCOMMAND_TEST_BOOL": "true", + }, + }, + { + name: "from env vars and config", + args: []string{"subcommand"}, + envVars: map[string]string{ + "TEST_STRING": "env-value", + }, + config: []byte(` +test-int: 123 +subcommand: + test-bool: true +`), + configFormat: "yaml", + }, + { + name: "from JSON config", + args: []string{"subcommand"}, + config: []byte(` + { + "test-string": "config-value", + "test-int": 123, + "subcommand": { + "test-bool": true + } + }`), + configFormat: "json", + }, + } + + var cmd *cobra.Command + var v *viper.Viper + + cobra.OnInitialize(func() { + configFilePath, err := cmd.Flags().GetString(configFlagName) + if err != nil { + cobra.CheckErr(err) + } + err = lib.LoadConfig(v, configFilePath) + assert.NoError(t, err) + err = lib.BindFlags(cmd, v, envVarPrefix) + assert.NoError(t, err) + }) + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + assertClearEnv(t) + for key, value := range tc.envVars { + err := setEnv(envVarPrefix+"_"+key, value) + assert.NoError(t, err) + } + defer clearEnvVars(t) + + var configFileName string + if tc.config != nil { + configFileName = writeTempFile(t, tc.config, tc.configFormat) + defer os.Remove(configFileName) + + tc.args = append(tc.args, "--"+configFlagName, configFileName) + } + + cmd = &cobra.Command{ + Use: "root", + } + testString := cmd.PersistentFlags().String("test-string", "", "Test string flag") + testInt := cmd.PersistentFlags().Int("test-int", 0, "Test int flag") + assert.NoError(t, cmd.MarkPersistentFlagRequired("test-string")) + cmd.PersistentFlags().String(configFlagName, "", "Config file name") + + var subcommandBool bool + var subCommandExecuted bool + subCmd := &cobra.Command{ + Use: "subcommand", + Run: func(cmd *cobra.Command, args []string) { + assert.NotEmpty(t, *testString) + assert.NotEmpty(t, *testInt) + assert.NotEmpty(t, subcommandBool) + subCommandExecuted = true + }, + } + subCmd.Flags().BoolVar(&subcommandBool, "test-bool", false, "Subcommand string flag") + cmd.AddCommand(subCmd) + + v = getViper() + + cmd.SetArgs(tc.args) + err := cmd.Execute() + assert.NoError(t, err) + + assert.True(t, subCommandExecuted) + subCommandExecuted = false + }) + } +} + +var envKeys []string + +func assertClearEnv(t *testing.T) { + assert.Len(t, envKeys, 0) +} + +func setEnv(key, value string) error { + envKeys = append(envKeys, key) + return os.Setenv(key, value) +} + +func clearEnvVars(t *testing.T) { + for len(envKeys) > 0 { + key := envKeys[0] + err := os.Unsetenv(key) + assert.NoError(t, err) + envKeys = envKeys[1:] + } +} + +func writeTempFile(t *testing.T, content []byte, fileExtension string) string { + file, err := os.CreateTemp("", "config-*."+fileExtension) + assert.NoError(t, err) + + _, err = file.Write([]byte(content)) + assert.NoError(t, err) + assert.NoError(t, file.Close()) + + return file.Name() +} + +func getViper() *viper.Viper { + v := viper.New() + v.SetEnvPrefix(envVarPrefix) + + return v +}