Merge pull request #105 from davidnewhall/dn2_more_fixes

Use reflection to get env var names.

.metadata.sh

@@ -11,7 +11,7 @@ HBREPO="golift/homebrew-mugs"
 MAINT="David Newhall II <david at sleepers dot pro>"
 VENDOR="Go Lift <code at golift dot io>"
 DESC="Polls a UniFi controller and exports metrics to InfluxDB"
-GOLANGCI_LINT_ARGS="--enable-all -D gochecknoglobals -e dupl -e G101"
+GOLANGCI_LINT_ARGS="--enable-all -D gochecknoglobals -D dupl -D lll -e G101 "
 # Example must exist at examples/$CONFIG_FILE.example
 CONFIG_FILE="up.conf"
 LICENSE="MIT"

unifipoller/config.go

@@ -1,11 +1,22 @@
 package unifipoller
 
 import (
+	"encoding/json"
+	"encoding/xml"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path"
+	"reflect"
+	"strconv"
+	"strings"
 	"time"
 
+	"github.com/BurntSushi/toml"
 	influx "github.com/influxdata/influxdb1-client/v2"
 	"github.com/spf13/pflag"
 	"golift.io/unifi"
+	"gopkg.in/yaml.v2"
 )
 
 // Version is injected by the Makefile
@@ -22,26 +33,9 @@ const (
 	defaultUnifURL  = "https://127.0.0.1:8443"
 )
 
-// These are environment variables that can be used to override configuration.
+// ENVConfigPrefix is the prefix appended to an env variable tag
-// Useful for Docker users.
+// name before retrieving the value from the OS.
-const (
+const ENVConfigPrefix = "UP_"
-	ENVConfigMode       = "UP_POLLING_MODE"
-	ENVConfigInfluxDB   = "UP_INFLUX_DB"
-	ENVConfigInfluxUser = "UP_INFLUX_USER"
-	ENVConfigInfluxPass = "UP_INFLUX_PASS"
-	ENVConfigInfluxURL  = "UP_INFLUX_URL"
-	ENVConfigUnifiUser  = "UP_UNIFI_USER"
-	ENVConfigUnifiPass  = "UP_UNIFI_PASS"
-	ENVConfigUnifiBase  = "UP_UNIFI_URL"
-	ENVConfigReAuth     = "UP_REAUTHENTICATE"
-	ENVConfigVerifySSL  = "UP_VERIFY_SSL"
-	ENVConfigCollectIDS = "UP_COLLECT_IDS"
-	ENVConfigQuiet      = "UP_QUIET_MODE"
-	ENVConfigDebug      = "UP_DEBUG_MODE"
-	ENVConfigInterval   = "UP_POLLING_INTERVAL"
-	ENVConfigMaxErrors  = "UP_MAX_ERRORS"
-	ENVConfigSites      = "UP_POLL_SITES"
-)
 
 // UnifiPoller contains the application startup data, and auth info for UniFi & Influx.
 type UnifiPoller struct {
@@ -73,23 +67,24 @@ type Metrics struct {
 
 // Config represents the data needed to poll a controller and report to influxdb.
 // This is all of the data stored in the config file.
+// Any with explicit defaults have _omitempty on json and toml tags.
 type Config struct {
-	MaxErrors  int      `json:"max_errors,_omitempty" toml:"max_errors,_omitempty" xml:"max_errors" yaml:"max_errors"`
+	MaxErrors  int      `json:"max_errors" toml:"max_errors" xml:"max_errors" yaml:"max_errors" env:"MAX_ERRORS"`
-	Interval   Duration `json:"interval,_omitempty" toml:"interval,_omitempty" xml:"interval" yaml:"interval"`
+	Interval   Duration `json:"interval,_omitempty" toml:"interval,_omitempty" xml:"interval" yaml:"interval" env:"POLLING_INTERVAL"`
-	Debug      bool     `json:"debug" toml:"debug" xml:"debug" yaml:"debug"`
+	Debug      bool     `json:"debug" toml:"debug" xml:"debug" yaml:"debug" env:"DEBUG_MODE"`
-	Quiet      bool     `json:"quiet,_omitempty" toml:"quiet,_omitempty" xml:"quiet" yaml:"quiet"`
+	Quiet      bool     `json:"quiet,_omitempty" toml:"quiet,_omitempty" xml:"quiet" yaml:"quiet" env:"QUIET_MODE"`
-	VerifySSL  bool     `json:"verify_ssl" toml:"verify_ssl" xml:"verify_ssl" yaml:"verify_ssl"`
+	VerifySSL  bool     `json:"verify_ssl" toml:"verify_ssl" xml:"verify_ssl" yaml:"verify_ssl" env:"VERIFY_SSL"`
-	CollectIDS bool     `json:"collect_ids" toml:"collect_ids" xml:"collect_ids" yaml:"collect_ids"`
+	CollectIDS bool     `json:"collect_ids" toml:"collect_ids" xml:"collect_ids" yaml:"collect_ids" env:"COLLECT_IDS"`
-	ReAuth     bool     `json:"reauthenticate" toml:"reauthenticate" xml:"reauthenticate" yaml:"reauthenticate"`
+	ReAuth     bool     `json:"reauthenticate" toml:"reauthenticate" xml:"reauthenticate" yaml:"reauthenticate" env:"REAUTHENTICATE"`
-	Mode       string   `json:"mode" toml:"mode" xml:"mode" yaml:"mode"`
+	Mode       string   `json:"mode" toml:"mode" xml:"mode" yaml:"mode" env:"POLLING_MODE"`
-	InfluxURL  string   `json:"influx_url,_omitempty" toml:"influx_url,_omitempty" xml:"influx_url" yaml:"influx_url"`
+	InfluxURL  string   `json:"influx_url,_omitempty" toml:"influx_url,_omitempty" xml:"influx_url" yaml:"influx_url" env:"INFLUX_URL"`
-	InfluxUser string   `json:"influx_user,_omitempty" toml:"influx_user,_omitempty" xml:"influx_user" yaml:"influx_user"`
+	InfluxUser string   `json:"influx_user,_omitempty" toml:"influx_user,_omitempty" xml:"influx_user" yaml:"influx_user" env:"INFLUX_USER"`
-	InfluxPass string   `json:"influx_pass,_omitempty" toml:"influx_pass,_omitempty" xml:"influx_pass" yaml:"influx_pass"`
+	InfluxPass string   `json:"influx_pass,_omitempty" toml:"influx_pass,_omitempty" xml:"influx_pass" yaml:"influx_pass" env:"INFLUX_PASS"`
-	InfluxDB   string   `json:"influx_db,_omitempty" toml:"influx_db,_omitempty" xml:"influx_db" yaml:"influx_db"`
+	InfluxDB   string   `json:"influx_db,_omitempty" toml:"influx_db,_omitempty" xml:"influx_db" yaml:"influx_db" env:"INFLUX_DB"`
-	UnifiUser  string   `json:"unifi_user,_omitempty" toml:"unifi_user,_omitempty" xml:"unifi_user" yaml:"unifi_user"`
+	UnifiUser  string   `json:"unifi_user,_omitempty" toml:"unifi_user,_omitempty" xml:"unifi_user" yaml:"unifi_user" env:"UNIFI_USER"`
-	UnifiPass  string   `json:"unifi_pass,_omitempty" toml:"unifi_pass,_omitempty" xml:"unifi_pass" yaml:"unifi_pass"`
+	UnifiPass  string   `json:"unifi_pass,_omitempty" toml:"unifi_pass,_omitempty" xml:"unifi_pass" yaml:"unifi_pass" env:"UNIFI_PASS"`
-	UnifiBase  string   `json:"unifi_url,_omitempty" toml:"unifi_url,_omitempty" xml:"unifi_url" yaml:"unifi_url"`
+	UnifiBase  string   `json:"unifi_url,_omitempty" toml:"unifi_url,_omitempty" xml:"unifi_url" yaml:"unifi_url" env:"UNIFI_URL"`
-	Sites      []string `json:"sites,_omitempty" toml:"sites,_omitempty" xml:"sites" yaml:"sites"`
+	Sites      []string `json:"sites,_omitempty" toml:"sites,_omitempty" xml:"sites" yaml:"sites" env:"POLL_SITES"`
 }
 
 // Duration is used to UnmarshalTOML into a time.Duration value.
@@ -100,3 +95,63 @@ func (d *Duration) UnmarshalText(data []byte) (err error) {
 	d.Duration, err = time.ParseDuration(string(data))
 	return
 }
+
+// ParseFile parses and returns our configuration data.
+func (c *Config) ParseFile(configFile string) error {
+	switch buf, err := ioutil.ReadFile(configFile); {
+	case err != nil:
+		return err
+	case strings.Contains(configFile, ".json"):
+		return json.Unmarshal(buf, c)
+	case strings.Contains(configFile, ".xml"):
+		return xml.Unmarshal(buf, c)
+	case strings.Contains(configFile, ".yaml"):
+		return yaml.Unmarshal(buf, c)
+	default:
+		return toml.Unmarshal(buf, c)
+	}
+}
+
+// ParseENV copies environment variables into configuration values.
+// This is useful for Docker users that find it easier to pass ENV variables
+// than a specific configuration file. Uses reflection to find struct tags.
+func (c *Config) ParseENV() error {
+	t := reflect.TypeOf(Config{}) // Get tag names from the Config struct.
+	// Loop each Config struct member; get reflect tag & env var value; update config.
+	for i := 0; i < t.NumField(); i++ {
+		tag := t.Field(i).Tag.Get("env")        // Get the ENV variable name from "env" struct tag
+		env := os.Getenv(ENVConfigPrefix + tag) // Then pull value from OS.
+		if tag == "" || env == "" {
+			continue // Skip if either are empty.
+		}
+
+		// Reflect and update the u.Config struct member at position i.
+		switch c := reflect.ValueOf(c).Elem().Field(i); c.Type().String() {
+		// Handle each member type appropriately (differently).
+		case "string":
+			// This is a reflect package method to update a struct member by index.
+			c.SetString(env)
+		case "int":
+			val, err := strconv.Atoi(env)
+			if err != nil {
+				return fmt.Errorf("%s: %v", tag, err)
+			}
+			c.Set(reflect.ValueOf(val))
+		case "[]string":
+			c.Set(reflect.ValueOf(strings.Split(env, ",")))
+		case path.Base(t.PkgPath()) + ".Duration":
+			val, err := time.ParseDuration(env)
+			if err != nil {
+				return fmt.Errorf("%s: %v", tag, err)
+			}
+			c.Set(reflect.ValueOf(Duration{val}))
+		case "bool":
+			val, err := strconv.ParseBool(env)
+			if err != nil {
+				return fmt.Errorf("%s: %v", tag, err)
+			}
+			c.SetBool(val)
+		}
+	}
+	return nil
+}

unifipoller/helpers.go

@@ -43,26 +43,3 @@ func (u *UnifiPoller) LogDebugf(m string, v ...interface{}) {
 func (u *UnifiPoller) LogErrorf(m string, v ...interface{}) {
 	_ = log.Output(2, fmt.Sprintf("[ERROR] "+m, v...))
 }
-
-// pick returns the first non empty string in a list.
-// used in a few places around this library.
-func pick(strings ...string) string {
-	for _, s := range strings {
-		if s != "" {
-			return s
-		}
-	}
-	return ""
-}
-
-// parseBool returns true/false if the string is "true" or "false", otherwise returns e value.
-func parseBool(s string, e bool) bool {
-	switch s {
-	case "true", "t":
-		return true
-	case "false", "f":
-		return false
-	default:
-		return e
-	}
-}

unifipoller/start.go

@@ -1,37 +1,51 @@
 package unifipoller
 
 import (
-	"encoding/json"
-	"encoding/xml"
 	"fmt"
-	"io/ioutil"
 	"log"
 	"os"
-	"strconv"
 	"strings"
 	"time"
 
-	"github.com/BurntSushi/toml"
 	influx "github.com/influxdata/influxdb1-client/v2"
 	"github.com/spf13/pflag"
 	"golift.io/unifi"
-	"gopkg.in/yaml.v2"
 )
 
 // Start begins the application from a CLI.
 // Parses flags, parses config and executes Run().
 func Start() error {
 	log.SetFlags(log.LstdFlags)
-	up := &UnifiPoller{Flag: &Flag{}}
+	up := &UnifiPoller{Flag: &Flag{},
+		Config: &Config{
+			// Preload our defaults.
+			InfluxURL:  defaultInfxURL,
+			InfluxUser: defaultInfxUser,
+			InfluxPass: defaultInfxPass,
+			InfluxDB:   defaultInfxDb,
+			UnifiUser:  defaultUnifUser,
+			UnifiPass:  os.Getenv("UNIFI_PASSWORD"), // deprecated name.
+			UnifiBase:  defaultUnifURL,
+			Interval:   Duration{defaultInterval},
+			Sites:      []string{"all"},
+		}}
 	up.Flag.Parse(os.Args[1:])
 	if up.Flag.ShowVer {
 		fmt.Printf("unifi-poller v%s\n", Version)
 		return nil // don't run anything else w/ version request.
 	}
-	if err := up.GetConfig(); err != nil {
+	if up.Flag.DumpJSON == "" { // do not print this when dumping JSON.
+		up.Logf("Loading Configuration File: %s", up.Flag.ConfigFile)
+	}
+	// Parse config file.
+	if err := up.Config.ParseFile(up.Flag.ConfigFile); err != nil {
 		up.Flag.Usage()
 		return err
 	}
+	// Update Config with ENV variable overrides.
+	if err := up.Config.ParseENV(); err != nil {
+		return err
+	}
 	return up.Run()
 }
 
@@ -43,71 +57,12 @@ func (f *Flag) Parse(args []string) {
 		f.PrintDefaults()
 	}
 	f.StringVarP(&f.DumpJSON, "dumpjson", "j", "",
-		"This debug option prints a json payload and exits. See man page for more.")
+		"This debug option prints a json payload and exits. See man page for more info.")
-	f.StringVarP(&f.ConfigFile, "config", "c", DefaultConfFile, "Poller Config File (TOML Format)")
+	f.StringVarP(&f.ConfigFile, "config", "c", DefaultConfFile, "Poller config file path.")
-	f.BoolVarP(&f.ShowVer, "version", "v", false, "Print the version and exit")
+	f.BoolVarP(&f.ShowVer, "version", "v", false, "Print the version and exit.")
 	_ = f.FlagSet.Parse(args)
 }
 
-// setEnvVarOptions copies environment variables into configuration values.
-// This is useful for Docker users that find it easier to pass ENV variables
-// that a specific configuration file.
-func (u *UnifiPoller) setEnvVarOptions() {
-	u.Config.Mode = pick(os.Getenv(ENVConfigMode), u.Config.Mode)
-	u.Config.InfluxDB = pick(os.Getenv(ENVConfigInfluxDB), u.Config.InfluxDB)
-	u.Config.InfluxUser = pick(os.Getenv(ENVConfigInfluxUser), u.Config.InfluxUser)
-	u.Config.InfluxPass = pick(os.Getenv(ENVConfigInfluxPass), u.Config.InfluxPass)
-	u.Config.InfluxURL = pick(os.Getenv(ENVConfigInfluxURL), u.Config.InfluxURL)
-	u.Config.UnifiUser = pick(os.Getenv(ENVConfigUnifiUser), u.Config.UnifiUser)
-	u.Config.UnifiPass = pick(os.Getenv(ENVConfigUnifiPass), u.Config.UnifiPass)
-	u.Config.UnifiBase = pick(os.Getenv(ENVConfigUnifiBase), u.Config.UnifiBase)
-	u.Config.ReAuth = parseBool(os.Getenv(ENVConfigReAuth), u.Config.ReAuth)
-	u.Config.VerifySSL = parseBool(os.Getenv(ENVConfigVerifySSL), u.Config.VerifySSL)
-	u.Config.CollectIDS = parseBool(os.Getenv(ENVConfigCollectIDS), u.Config.CollectIDS)
-	u.Config.Quiet = parseBool(os.Getenv(ENVConfigQuiet), u.Config.Quiet)
-	u.Config.Debug = parseBool(os.Getenv(ENVConfigDebug), u.Config.Debug)
-	if e := os.Getenv(ENVConfigInterval); e != "" {
-		_ = u.Config.Interval.UnmarshalText([]byte(e))
-	}
-	if e := os.Getenv(ENVConfigMaxErrors); e != "" {
-		u.Config.MaxErrors, _ = strconv.Atoi(e)
-	}
-	if e := os.Getenv(ENVConfigSites); e != "" {
-		u.Config.Sites = strings.Split(e, ",")
-	}
-}
-
-// GetConfig parses and returns our configuration data.
-func (u *UnifiPoller) GetConfig() error {
-	// Preload our defaults.
-	u.Config = &Config{
-		InfluxURL:  defaultInfxURL,
-		InfluxUser: defaultInfxUser,
-		InfluxPass: defaultInfxPass,
-		InfluxDB:   defaultInfxDb,
-		UnifiUser:  defaultUnifUser,
-		UnifiPass:  os.Getenv("UNIFI_PASSWORD"), // deprecated name.
-		UnifiBase:  defaultUnifURL,
-		Interval:   Duration{defaultInterval},
-		Sites:      []string{"default"},
-		Quiet:      u.Flag.DumpJSON != "", //s uppress the following u.Logf line.
-	}
-	u.Logf("Loading Configuration File: %s", u.Flag.ConfigFile)
-	defer u.setEnvVarOptions() // Set env variable overrides when done here.
-	switch buf, err := ioutil.ReadFile(u.Flag.ConfigFile); {
-	case err != nil:
-		return err
-	case strings.Contains(u.Flag.ConfigFile, ".json"):
-		return json.Unmarshal(buf, u.Config)
-	case strings.Contains(u.Flag.ConfigFile, ".xml"):
-		return xml.Unmarshal(buf, u.Config)
-	case strings.Contains(u.Flag.ConfigFile, ".yaml"):
-		return yaml.Unmarshal(buf, u.Config)
-	default:
-		return toml.Unmarshal(buf, u.Config)
-	}
-}
-
 // Run invokes all the application logic and routines.
 func (u *UnifiPoller) Run() (err error) {
 	if u.Flag.DumpJSON != "" {
@@ -164,8 +119,5 @@ func (u *UnifiPoller) GetUnifi() (err error) {
 	if err != nil {
 		return fmt.Errorf("unifi controller: %v", err)
 	}
-	if err := u.CheckSites(); err != nil {
+	return u.CheckSites()
-		return err
-	}
-	return nil
 }

unifipoller/unifi.go

@@ -62,7 +62,8 @@ func (u *UnifiPoller) PollController() error {
 			_ = u.CollectAndReport()
 		}
 		if u.Config.MaxErrors >= 0 && u.errorCount > u.Config.MaxErrors {
-			return fmt.Errorf("reached maximum error count, stopping poller (%d > %d)", u.errorCount, u.Config.MaxErrors)
+			return fmt.Errorf("reached maximum error count, stopping poller (%d > %d)",
+				u.errorCount, u.Config.MaxErrors)
 		}
 	}
 	return nil