vendor/github.com/tendermint/tmlibs/cli/setup.go

   1 package cli
   2
   3 import (
   4         "fmt"
   5         "os"
   6         "strings"
   7
   8         "github.com/pkg/errors"
   9         "github.com/spf13/cobra"
  10         "github.com/spf13/viper"
  11 )
  12
  13 const (
  14         RootFlag     = "root"
  15         HomeFlag     = "home"
  16         TraceFlag    = "trace"
  17         OutputFlag   = "output"
  18         EncodingFlag = "encoding"
  19 )
  20
  21 // Executable is the minimal interface to *corba.Command, so we can
  22 // wrap if desired before the test
  23 type Executable interface {
  24         Execute() error
  25 }
  26
  27 // PrepareBaseCmd is meant for tendermint and other servers
  28 func PrepareBaseCmd(cmd *cobra.Command, envPrefix, defautRoot string) Executor {
  29         cobra.OnInitialize(func() { initEnv(envPrefix) })
  30         cmd.PersistentFlags().StringP(RootFlag, "r", defautRoot, "DEPRECATED. Use --home")
  31         // -h is already reserved for --help as part of the cobra framework
  32         // do you want to try something else??
  33         // also, default must be empty, so we can detect this unset and fall back
  34         // to --root / TM_ROOT / TMROOT
  35         cmd.PersistentFlags().String(HomeFlag, "", "root directory for config and data")
  36         cmd.PersistentFlags().Bool(TraceFlag, false, "print out full stack trace on errors")
  37         cmd.PersistentPreRunE = concatCobraCmdFuncs(bindFlagsLoadViper, cmd.PersistentPreRunE)
  38         return Executor{cmd, os.Exit}
  39 }
  40
  41 // initEnv sets to use ENV variables if set.
  42 func initEnv(prefix string) {
  43         copyEnvVars(prefix)
  44
  45         // env variables with TM prefix (eg. TM_ROOT)
  46         viper.SetEnvPrefix(prefix)
  47         viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_", "-", "_"))
  48         viper.AutomaticEnv()
  49 }
  50
  51 // This copies all variables like TMROOT to TM_ROOT,
  52 // so we can support both formats for the user
  53 func copyEnvVars(prefix string) {
  54         prefix = strings.ToUpper(prefix)
  55         ps := prefix + "_"
  56         for _, e := range os.Environ() {
  57                 kv := strings.SplitN(e, "=", 2)
  58                 if len(kv) == 2 {
  59                         k, v := kv[0], kv[1]
  60                         if strings.HasPrefix(k, prefix) && !strings.HasPrefix(k, ps) {
  61                                 k2 := strings.Replace(k, prefix, ps, 1)
  62                                 os.Setenv(k2, v)
  63                         }
  64                 }
  65         }
  66 }
  67
  68 // Executor wraps the cobra Command with a nicer Execute method
  69 type Executor struct {
  70         *cobra.Command
  71         Exit func(int) // this is os.Exit by default, override in tests
  72 }
  73
  74 type ExitCoder interface {
  75         ExitCode() int
  76 }
  77
  78 // execute adds all child commands to the root command sets flags appropriately.
  79 // This is called by main.main(). It only needs to happen once to the rootCmd.
  80 func (e Executor) Execute() error {
  81         e.SilenceUsage = true
  82         e.SilenceErrors = true
  83         err := e.Command.Execute()
  84         if err != nil {
  85                 if viper.GetBool(TraceFlag) {
  86                         fmt.Fprintf(os.Stderr, "ERROR: %+v\n", err)
  87                 } else {
  88                         fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
  89                 }
  90
  91                 // return error code 1 by default, can override it with a special error type
  92                 exitCode := 1
  93                 if ec, ok := err.(ExitCoder); ok {
  94                         exitCode = ec.ExitCode()
  95                 }
  96                 e.Exit(exitCode)
  97         }
  98         return err
  99 }
 100
 101 type cobraCmdFunc func(cmd *cobra.Command, args []string) error
 102
 103 // Returns a single function that calls each argument function in sequence
 104 // RunE, PreRunE, PersistentPreRunE, etc. all have this same signature
 105 func concatCobraCmdFuncs(fs ...cobraCmdFunc) cobraCmdFunc {
 106         return func(cmd *cobra.Command, args []string) error {
 107                 for _, f := range fs {
 108                         if f != nil {
 109                                 if err := f(cmd, args); err != nil {
 110                                         return err
 111                                 }
 112                         }
 113                 }
 114                 return nil
 115         }
 116 }
 117
 118 // Bind all flags and read the config into viper
 119 func bindFlagsLoadViper(cmd *cobra.Command, args []string) error {
 120         // cmd.Flags() includes flags from this command and all persistent flags from the parent
 121         if err := viper.BindPFlags(cmd.Flags()); err != nil {
 122                 return err
 123         }
 124
 125         // rootDir is command line flag, env variable, or default $HOME/.tlc
 126         // NOTE: we support both --root and --home for now, but eventually only --home
 127         // Also ensure we set the correct rootDir under HomeFlag so we dont need to
 128         // repeat this logic elsewhere.
 129         rootDir := viper.GetString(HomeFlag)
 130         if rootDir == "" {
 131                 rootDir = viper.GetString(RootFlag)
 132                 viper.Set(HomeFlag, rootDir)
 133         }
 134         viper.SetConfigName("config") // name of config file (without extension)
 135         viper.AddConfigPath(rootDir)  // search root directory
 136
 137         // If a config file is found, read it in.
 138         if err := viper.ReadInConfig(); err == nil {
 139                 // stderr, so if we redirect output to json file, this doesn't appear
 140                 // fmt.Fprintln(os.Stderr, "Using config file:", viper.ConfigFileUsed())
 141         } else if _, ok := err.(viper.ConfigFileNotFoundError); !ok {
 142                 // we ignore not found error, only parse error
 143                 // stderr, so if we redirect output to json file, this doesn't appear
 144                 fmt.Fprintf(os.Stderr, "%#v", err)
 145         }
 146         return nil
 147 }
 148
 149 func validateOutput(cmd *cobra.Command, args []string) error {
 150         // validate output format
 151         output := viper.GetString(OutputFlag)
 152         switch output {
 153         case "text", "json":
 154         default:
 155                 return errors.Errorf("Unsupported output format: %s", output)
 156         }
 157         return nil
 158 }