reflector.go/cmd/root.go

package cmd

import (
	"encoding/json"
	"io/ioutil"
	"os"

	"github.com/lbryio/lbry.go/errors"
	"github.com/lbryio/lbry.go/util"
	"github.com/lbryio/reflector.go/dht"

	"github.com/johntdyer/slackrus"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
)

// Config is the base configuration for Prism when no sub commands are used.
type Config struct {
	AwsID        string `json:"aws_id"`
	AwsSecret    string `json:"aws_secret"`
	BucketRegion string `json:"bucket_region"`
	BucketName   string `json:"bucket_name"`
	DBConn       string `json:"db_conn"`
	SlackHookURL string `json:"slack_hook_url"`
}

var verbose []string

const (
	verboseAll        = "all"
	verboseDHT        = "dht"
	verboseNodeFinder = "node_finder"
)

var conf string
var globalConfig Config

var rootCmd = &cobra.Command{
	Use:              "prism",
	Short:            "Prism is a single entry point application with multiple sub modules which can be leveraged individually or together",
	PersistentPreRun: preRun,
	// Uncomment the following line if your bare application
	// has an action associated with it:
	//	Run: func(cmd *cobra.Command, args []string) { },
}

func init() {
	rootCmd.PersistentFlags().StringSliceVarP(&verbose, "verbose", "v", []string{}, "Verbose logging for specific components")
	rootCmd.PersistentFlags().StringVar(&conf, "conf", "config.json", "Path to config")
}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	err := rootCmd.Execute()
	if err != nil {
		logrus.Errorln(err)
		os.Exit(1)
	}
}

func preRun(cmd *cobra.Command, args []string) {
	debugLogger := logrus.New()
	debugLogger.SetLevel(logrus.DebugLevel)
	debugLogger.SetOutput(os.Stderr)

	if util.InSlice(verboseAll, verbose) {
		logrus.SetLevel(logrus.DebugLevel)
		verbose = []string{verboseDHT, verboseNodeFinder}
	}

	for _, debugType := range verbose {
		switch debugType {
		case verboseDHT:
			dht.UseLogger(debugLogger)
		case verboseNodeFinder:
			dht.NodeFinderUseLogger(debugLogger)
		}
	}

	var err error
	if conf == "" {
		logrus.Errorln("--conf flag required")
		os.Exit(1)
	} else if conf != "none" {
		globalConfig, err = loadConfig(conf)
		if err != nil {
			logrus.Error(err)
			os.Exit(1)
		}
	}

	if globalConfig.SlackHookURL != "" {
		hook := &slackrus.SlackrusHook{
			HookURL:        globalConfig.SlackHookURL,
			AcceptedLevels: slackrus.LevelThreshold(logrus.InfoLevel),
			Channel:        "#reflector-logs",
			//IconEmoji:      ":ghost:",
			//Username:       "reflector.go",
		}
		//logrus.SetFormatter(&logrus.JSONFormatter{})
		logrus.AddHook(hook)
		debugLogger.AddHook(hook)
	}
}

func checkErr(err error) {
	if err != nil {
		panic(err)
	}
}

func argFuncs(funcs ...cobra.PositionalArgs) cobra.PositionalArgs {
	return func(cmd *cobra.Command, args []string) error {
		for _, f := range funcs {
			err := f(cmd, args)
			if err != nil {
				return err
			}
		}
		return nil
	}
}

func loadConfig(path string) (Config, error) {
	var c Config

	raw, err := ioutil.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return c, errors.Err("config file not found")
		}
		return c, err
	}

	err = json.Unmarshal(raw, &c)
	return c, err
}