makeshift/cmd/root.go

package cmd

import (
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"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",
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		var (
			logFile, _    = cmd.Flags().GetString("log-file")
			configPath, _ = cmd.Flags().GetString("config")
			err           error
		)

		// initialize the logger
		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
		return initConfig(cmd, configPath)
	},
	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")
		setenv(cmd, "config", "MAKESHIFT_CONFIG_FILE")
		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")
			os.Exit(1)
		}
	},
}

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 (can be set with MAKESHIFT_LOG_LEVEL)")
	rootCmd.PersistentFlags().String("log-file", "", "Set the log file path (can be set with MAKESHIFT_LOG_FILE)")
	rootCmd.PersistentFlags().StringP("config", "c", "", "Set the config file path (can be set with MAKESHIFT_CONFIG_FILE)")
}

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 initConfig(cmd *cobra.Command, path string) error {

	// Dissect the path to separate config name from its directory
	var (
		isFlagSet = cmd.Flags().Changed("config")
		filename  = filepath.Base(path)
		ext       = filepath.Ext(filename)
		directory = filepath.Dir(path)
		v         = viper.New()
	)

	// The 'config' flag not set, so don't continue
	if !isFlagSet {
		return nil
	}

	// Only use specified YAML file from --config or -c flag
	v.SetConfigName(strings.TrimSuffix(filename, ext))
	v.SetConfigType("yaml")
	v.AddConfigPath(directory)

	// Attempt to read the config file. Return an error if we cannot parse
	// the config file or if it is not found.
	if err := v.ReadInConfig(); err != nil {
		if isFlagSet {
			// It's okay if there isn't a config file when no path is provided
			// if _, ok := err.(viper.ConfigFileNotFoundError); !ok {
			// 	return err
			// }
			switch err.(type) {
			case viper.ConfigFileNotFoundError:
				log.Error().
					Err(err).
					Str("path", path).
					Msg("failed to read config")
				os.Exit(1)
			default:
				log.Error().Err(err).Msg("failed to read config")
				os.Exit(1)
			}
		}
	}

	// 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
}

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 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))
		}
	})
}