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