package cmd import ( "fmt" "net/http" "os" "strings" logger "git.towk2.me/towk/makeshift/pkg/log" "github.com/rs/zerolog/log" "github.com/spf13/cobra" "github.com/spf13/pflag" "github.com/spf13/viper" ) var ( loglevel logger.LogLevel = logger.INFO ) var rootCmd = cobra.Command{ Use: "makeshift", Short: "Extensible file cobbler", PersistentPreRun: func(cmd *cobra.Command, args []string) { var ( logFile string err error ) // initialize the logger logFile, _ = cmd.Flags().GetString("log-file") err = logger.InitWithLogLevel(loglevel, logFile) if err != nil { log.Error().Err(err).Msg("failed to initialize logger") os.Exit(1) } // You can bind cobra and viper in a few locations, but PersistencePreRunE on the root command works well err = initializeConfig(cmd) }, Run: func(cmd *cobra.Command, args []string) { // try and set flags using env vars setenv(cmd, "log-file", "MAKESHIFT_LOG_FILE") setenv(cmd, "log-level", "MAKESHIFT_LOG_LEVEL") if len(args) == 0 { err := cmd.Help() if err != nil { log.Error().Err(err).Msg("failed to print help") } os.Exit(0) } }, PostRun: func(cmd *cobra.Command, args []string) { log.Debug().Msg("closing log file") err := logger.LogFile.Close() if err != nil { log.Error().Err(err).Msg("failed to close log file") } }, } func Execute() { // run the main program if err := rootCmd.Execute(); err != nil { fmt.Fprintln(os.Stderr, err) os.Exit(1) } } func init() { cobra.OnInitialize( initLogger, ) // initialize the config a single time rootCmd.PersistentFlags().VarP(&loglevel, "log-level", "l", "Set the log level output") rootCmd.PersistentFlags().String("log-file", "", "Set the log file path (can be set with MAKESHIFT_LOG_FILE)") } func setenv(cmd *cobra.Command, varname string, envvar string) { if cmd.Flags().Changed(varname) { return } val := os.Getenv(envvar) if val != "" { cmd.Flags().Set(varname, val) } } func setenvp(cmd *cobra.Command, varname string, envvar string) { if cmd.Flags().Changed(varname) { return } val := os.Getenv(envvar) if val != "" { cmd.PersistentFlags().Set(varname, val) } } func initLogger() { // initialize the logger logfile, _ := rootCmd.PersistentFlags().GetString("log-file") err := logger.InitWithLogLevel(loglevel, logfile) if err != nil { log.Error().Err(err).Msg("failed to initialize logger") os.Exit(1) } } func handleResponseError(res *http.Response, host, query string, err error) { if err != nil { log.Error().Err(err). Str("host", host). Str("query", query). Msg("failed to make request") os.Exit(1) } if res.StatusCode != http.StatusOK { log.Error(). Any("status", map[string]any{ "code": res.StatusCode, "message": res.Status, }). Str("host", host). Msg("response returned bad status") os.Exit(1) } } // helper to write downloaded files func writeFiles(path string, body []byte) { var err = os.WriteFile(path, body, 0o755) if err != nil { log.Error().Err(err).Msg("failed to write file(s) from download") os.Exit(1) } } // Bind each cobra flag to its associated viper configuration (config file and environment variable) func bindFlags(cmd *cobra.Command, v *viper.Viper) { cmd.Flags().VisitAll(func(f *pflag.Flag) { // Determine the naming convention of the flags when represented in the config file configName := f.Name // Apply the viper config value to the flag when the flag is not set and viper has a value if !f.Changed && v.IsSet(configName) { val := v.Get(configName) cmd.Flags().Set(f.Name, fmt.Sprintf("%v", val)) } }) } func initializeConfig(cmd *cobra.Command) error { v := viper.New() // Set the base name of the config file, without the file extension. v.SetConfigName("makeshift") // Set as many paths as you like where viper should look for the // config file. We are only looking in the current working directory. v.AddConfigPath(".") // Attempt to read the config file, gracefully ignoring errors // caused by a config file not being found. Return an error // if we cannot parse the config file. if err := v.ReadInConfig(); err != nil { // It's okay if there isn't a config file if _, ok := err.(viper.ConfigFileNotFoundError); !ok { return err } } // When we bind flags to environment variables expect that the // environment variables are prefixed, e.g. a flag like --number // binds to an environment variable STING_NUMBER. This helps // avoid conflicts. v.SetEnvPrefix("MAKESHIFT") // Environment variables can't have dashes in them, so bind them to their equivalent // keys with underscores, e.g. --favorite-color to STING_FAVORITE_COLOR v.SetEnvKeyReplacer(strings.NewReplacer("-", "_")) // Bind to environment variables // Works great for simple config names, but needs help for names // like --favorite-color which we fix in the bindFlags function v.AutomaticEnv() // Bind the current command's flags to viper bindFlags(cmd, v) return nil }