bottin/config.go

package main

import (
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const (
	ViperAPITLSEnabled       string = "api.tls.enabled"
	FlagAPITLSEnabled        string = "api-tls-enabled"
	DefaultAPITLSEnabled     bool   = false
	DescriptionAPITLSEnabled string = "Whether to use TLS or not. Requires certificate and private key files."

	ViperAPITLSCertificateFile       string = "api.tls.certificate_file"
	FlagAPITLSCertificateFile        string = "api-tls-certificate-file"
	DefaultAPITLSCertificateFile     string = ""
	DescriptionAPITLSCertificateFile string = "Path to TLS certificate file"

	ViperAPITLSPrivateKeyFile       string = "api.tls.private_key_file"
	FlagAPITLSPrivateKeyFile        string = "api-tls-private-key-file"
	DefaultAPITLSPrivateKeyFile     string = ""
	DescriptionAPITLSPrivateKeyFile string = "Path to TLS private key file"

	ViperAPIPort       string = "api.port"
	FlagAPIPort        string = "api-port"
	DefaultAPIPort     int    = 1312
	DescriptionAPIPort string = "API server port"

	ViperAPIKey       string = "api.key"
	FlagAPIKey        string = "api-key"
	DefaultAPIKey     string = "bottin"
	DescriptionAPIKey string = "API server key. Leave empty for no key auth (not recommended)"

	ViperDBDatabase       string = "db.database"
	FlagDBDatabase        string = "db-database"
	DefaultDBDatabase     string = "bottin"
	DescriptionDBDatabase string = "Postgres database"

	ViperDBSSLMode       string = "db.sslmode"
	FlagDBSSLMode        string = "db-sslmode"
	DefaultDBSSLMode     string = "prefer"
	DescriptionDBSSLMode string = "Postgres sslmode"

	ViperDBHost       string = "db.host"
	FlagDBHost        string = "db-host"
	DefaultDBHost     string = "db"
	DescriptionDBHost string = "Postgres host"

	ViperDBPassword       string = "db.password"
	FlagDBPassword        string = "db-password"
	DefaultDBPassword     string = "bottin"
	DescriptionDBPassword string = "Postgres password"

	ViperDBPort       string = "db.port"
	FlagDBPort        string = "db-port"
	DefaultDBPort     int    = 5432
	DescriptionDBPort string = "Postgres port"

	ViperDBUser       string = "db.user"
	FlagDBUser        string = "db-user"
	DefaultDBUser     string = "bottin"
	DescriptionDBUser string = "Postgres user"

	ViperWebUser       string = "web.user"
	FlagWebUser        string = "web-user"
	DefaultWebUser     string = "bottin"
	DescriptionWebUser string = "Web client basic auth user"

	ViperWebPassword       string = "web.password"
	FlagWebPassword        string = "web-password"
	DefaultWebPassword     string = "bottin"
	DescriptionWebPassword string = "Web client basic auth password"

	ViperWebPort       string = "web.port"
	FlagWebPort        string = "web-port"
	DefaultWebPort     int    = 2312
	DescriptionWebPort string = "Web client port"

	ViperWebAPIHost       string = "web.api.host"
	FlagWebAPIHost        string = "web-api-host"
	DefaultWebAPIHost     string = "api"
	DescriptionWebAPIHost string = "Target API server host"

	ViperWebAPIKey       string = "web.api.key"
	FlagWebAPIKey        string = "web-api-key"
	DefaultWebAPIKey     string = "bottin"
	DescriptionWebAPIKey string = "Target API server key"

	ViperWebAPIPort       string = "web.api.port"
	FlagWebAPIPort        string = "web-api-port"
	DefaultWebAPIPort     int    = 1312
	DescriptionWebAPIPort string = "Target API server port"

	ViperWebAPIProtocol       string = "web.api.protocol"
	FlagWebAPIProtocol        string = "web-api-protocol"
	DefaultWebAPIProtocol     string = "http"
	DescriptionWebAPIProtocol string = "Target API server protocol (http/https)"
)

type Config struct {
	API struct {
		TLS struct {
			Enabled bool `yaml:"enabled"`

			// Path to file containing TLS certificate
			CertificateFile string `yaml:"certificate_file"`

			// Path to file containing TLS private key
			PrivateKeyFile string `yaml:"private_key_file"`
		}
		Port int    `yaml:"port"`
		Key  string `yaml:"key"`
	} `yaml:"api"`
	DB struct {
		Database string `yaml:"database"`
		Host     string `yaml:"host"`
		SSLMode  string `yaml:"sslmode"`
		Password string `yaml:"password"`
		Port     int    `yaml:"port"`
		User     string `yaml:"user"`
	} `yaml:"db"`
	Web struct {
		User     string `yaml:"user"`
		Password string `yaml:"password"`
		Port     int    `yaml:"port"`
		API      struct {
			Host     string `yaml:"host"`
			Key      string `yaml:"key"`
			Port     int    `yaml:"port"`
			Protocol string `yaml:"protocol"`
		} `yaml:"api"`
	} `yaml:"web"`
}

// DefaultConfig returns a Config filled with the default values from the
// `Default*` constants defined in this file.
func DefaultConfig() (cfg Config) {
	cfg.API.TLS.Enabled = DefaultAPITLSEnabled
	cfg.API.TLS.CertificateFile = DefaultAPITLSCertificateFile
	cfg.API.TLS.PrivateKeyFile = DefaultAPITLSPrivateKeyFile
	cfg.API.Port = DefaultAPIPort
	cfg.API.Key = DefaultAPIKey
	cfg.DB.Database = DefaultDBDatabase
	cfg.DB.Host = DefaultDBHost
	cfg.DB.SSLMode = DefaultDBSSLMode
	cfg.DB.Password = DefaultDBPassword
	cfg.DB.Port = DefaultDBPort
	cfg.DB.User = DefaultDBUser
	cfg.Web.User = DefaultWebUser
	cfg.Web.Password = DefaultWebPassword
	cfg.Web.Port = DefaultWebPort
	cfg.Web.API.Host = DefaultWebAPIHost
	cfg.Web.API.Key = DefaultWebAPIKey
	cfg.Web.API.Port = DefaultWebAPIPort
	cfg.Web.API.Protocol = DefaultWebAPIProtocol
	return
}

func init() {
	// rootCmd

	cobra.OnInitialize(initConfig)

	rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.bottin.yaml)")

	// apiCmd

	rootCmd.AddCommand(apiCmd)

	// api.tls.enabled
	apiCmd.Flags().Bool(FlagAPITLSEnabled, DefaultAPITLSEnabled, DescriptionAPITLSEnabled)
	if err := viper.BindPFlag(ViperAPITLSEnabled, apiCmd.Flags().Lookup(FlagAPITLSEnabled)); err != nil {
		log.Fatal(err)
	}

	// api.tls.certificate_file
	apiCmd.Flags().String(FlagAPITLSCertificateFile, DefaultAPITLSCertificateFile, DescriptionAPITLSCertificateFile)
	if err := viper.BindPFlag(ViperAPITLSCertificateFile, apiCmd.Flags().Lookup(FlagAPITLSCertificateFile)); err != nil {
		log.Fatal(err)
	}

	// api.tls.private_key_file
	apiCmd.Flags().String(FlagAPITLSPrivateKeyFile, DefaultAPITLSPrivateKeyFile, DescriptionAPITLSPrivateKeyFile)
	if err := viper.BindPFlag(ViperAPITLSPrivateKeyFile, apiCmd.Flags().Lookup(FlagAPITLSPrivateKeyFile)); err != nil {
		log.Fatal(err)
	}

	// api.key
	apiCmd.Flags().String(FlagAPIKey, DefaultAPIKey, DescriptionAPIKey)
	if err := viper.BindPFlag(ViperAPIKey, apiCmd.Flags().Lookup(FlagAPIKey)); err != nil {
		log.Fatal(err)
	}

	// api.port
	apiCmd.Flags().Int(FlagAPIPort, DefaultAPIPort, DescriptionAPIPort)
	if err := viper.BindPFlag(ViperAPIPort, apiCmd.Flags().Lookup(FlagAPIPort)); err != nil {
		log.Fatal(err)
	}

	// db.database
	apiCmd.Flags().String(FlagDBDatabase, DefaultDBDatabase, DescriptionDBDatabase)
	if err := viper.BindPFlag(ViperDBDatabase, apiCmd.Flags().Lookup(FlagDBDatabase)); err != nil {
		log.Fatal(err)
	}

	// db.sslmode
	apiCmd.Flags().String(FlagDBSSLMode, DefaultDBSSLMode, DescriptionDBSSLMode)
	if err := viper.BindPFlag(ViperDBSSLMode, apiCmd.Flags().Lookup(FlagDBSSLMode)); err != nil {
		log.Fatal(err)
	}

	// db.host
	apiCmd.Flags().String(FlagDBHost, DefaultDBHost, DescriptionDBHost)
	if err := viper.BindPFlag(ViperDBHost, apiCmd.Flags().Lookup(FlagDBHost)); err != nil {
		log.Fatal(err)
	}

	// db.password
	apiCmd.Flags().String(FlagDBPassword, DefaultDBPassword, DescriptionDBPassword)
	if err := viper.BindPFlag(ViperDBPassword, apiCmd.Flags().Lookup(FlagDBPassword)); err != nil {
		log.Fatal(err)
	}

	// db.port
	apiCmd.Flags().Int(FlagDBPort, DefaultDBPort, DescriptionDBPort)
	if err := viper.BindPFlag(ViperDBPort, apiCmd.Flags().Lookup(FlagDBPort)); err != nil {
		log.Fatal(err)
	}

	// db.user
	apiCmd.Flags().String(FlagDBUser, DefaultDBUser, DescriptionDBUser)
	if err := viper.BindPFlag(ViperDBUser, apiCmd.Flags().Lookup(FlagDBUser)); err != nil {
		log.Fatal(err)
	}

	// WebCmd
	rootCmd.AddCommand(webCmd)

	// web.api.host
	webCmd.Flags().String(FlagWebAPIHost, DefaultWebAPIHost, DescriptionWebAPIHost)
	if err := viper.BindPFlag(ViperWebAPIHost, webCmd.Flags().Lookup(FlagWebAPIHost)); err != nil {
		log.Fatal(err)
	}

	// web.api.key
	webCmd.Flags().String(FlagWebAPIKey, DefaultWebAPIKey, DescriptionWebAPIKey)
	if err := viper.BindPFlag(ViperWebAPIKey, webCmd.Flags().Lookup(FlagWebAPIKey)); err != nil {
		log.Fatal(err)
	}

	// web.api.protocol
	webCmd.Flags().String(FlagWebAPIProtocol, DefaultWebAPIProtocol, DescriptionWebAPIProtocol)
	if err := viper.BindPFlag(ViperWebAPIProtocol, webCmd.Flags().Lookup(FlagWebAPIProtocol)); err != nil {
		log.Fatal(err)
	}

	// web.api.port
	webCmd.Flags().Int(FlagWebAPIPort, DefaultWebAPIPort, DescriptionWebAPIPort)
	if err := viper.BindPFlag(ViperWebAPIPort, webCmd.Flags().Lookup(FlagWebAPIPort)); err != nil {
		log.Fatal(err)
	}

	// web.password
	webCmd.Flags().String(FlagWebPassword, DefaultWebPassword, DescriptionWebPassword)
	if err := viper.BindPFlag(ViperWebPassword, webCmd.Flags().Lookup(FlagWebPassword)); err != nil {
		log.Fatal(err)
	}

	// web.port
	webCmd.Flags().Int(FlagWebPort, DefaultWebPort, DescriptionWebPort)
	if err := viper.BindPFlag(ViperWebPort, webCmd.Flags().Lookup(FlagWebPort)); err != nil {
		log.Fatal(err)
	}

	// web.user
	webCmd.Flags().String(FlagWebUser, DefaultWebUser, DescriptionWebUser)
	if err := viper.BindPFlag(ViperWebUser, webCmd.Flags().Lookup(FlagWebUser)); err != nil {
		log.Fatal(err)
	}
}

var cfgFile string

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if cfgFile != "" {
		// Use config file from the flag.
		viper.SetConfigFile(cfgFile)
	} else {
		// Find home directory.
		home, err := os.UserHomeDir()
		cobra.CheckErr(err)

		// Search config in home directory with name ".bottin" (without extension).
		viper.AddConfigPath(home)
		viper.SetConfigType("yaml")
		viper.SetConfigName(".bottin")
	}

	viper.SetEnvPrefix("BOTTIN")
	viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
	viper.AutomaticEnv() // read in environment variables that match

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err == nil {
		fmt.Fprintln(os.Stderr, "Using config file:", viper.ConfigFileUsed())
	}
}