You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

178 lines
5.4 KiB

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