Remove struct inheritance.

This commit is contained in:
David Newhall II 2019-08-25 19:01:29 -07:00
parent 13b417c52f
commit 2a355c0677
5 changed files with 79 additions and 71 deletions

View File

@ -24,15 +24,20 @@ const (
// UnifiPoller contains the application startup data, and auth info for UniFi & Influx.
type UnifiPoller struct {
Influx influx.Client
Unifi *unifi.Unifi
Flag *Flag
Config *Config
errorCount int
LastCheck time.Time
}
// Flag represents the CLI args available and their settings.
type Flag struct {
ConfigFile string
DumpJSON string
ShowVer bool
Flag *pflag.FlagSet
errorCount int
LastCheck time.Time
influx.Client
*unifi.Unifi
*Config
*pflag.FlagSet
}
// Metrics contains all the data from the controller and an influx endpoint to send it to.

View File

@ -10,17 +10,17 @@ import (
// DumpJSONPayload prints raw json from the UniFi Controller.
func (u *UnifiPoller) DumpJSONPayload() (err error) {
u.Quiet = true
u.Config.Quiet = true
u.Unifi, err = unifi.NewUnifi(&unifi.Config{
User: u.UnifiUser,
Pass: u.UnifiPass,
URL: u.UnifiBase,
VerifySSL: u.VerifySSL,
User: u.Config.UnifiUser,
Pass: u.Config.UnifiPass,
URL: u.Config.UnifiBase,
VerifySSL: u.Config.VerifySSL,
})
if err != nil {
return err
}
fmt.Fprintln(os.Stderr, "[INFO] Authenticated to UniFi Controller @", u.UnifiBase, "as user", u.UnifiUser)
fmt.Fprintln(os.Stderr, "[INFO] Authenticated to UniFi Controller @", u.Config.UnifiBase, "as user", u.Config.UnifiUser)
if err := u.CheckSites(); err != nil {
return err
}
@ -30,12 +30,12 @@ func (u *UnifiPoller) DumpJSONPayload() (err error) {
switch sites, err := u.GetFilteredSites(); {
case err != nil:
return err
case StringInSlice(u.DumpJSON, []string{"d", "device", "devices"}):
case StringInSlice(u.Flag.DumpJSON, []string{"d", "device", "devices"}):
return u.dumpSitesJSON(unifi.DevicePath, "Devices", sites)
case StringInSlice(u.DumpJSON, []string{"client", "clients", "c"}):
case StringInSlice(u.Flag.DumpJSON, []string{"client", "clients", "c"}):
return u.dumpSitesJSON(unifi.ClientPath, "Clients", sites)
case strings.HasPrefix(u.DumpJSON, "other "):
apiPath := strings.SplitN(u.DumpJSON, " ", 2)[1]
case strings.HasPrefix(u.Flag.DumpJSON, "other "):
apiPath := strings.SplitN(u.Flag.DumpJSON, " ", 2)[1]
_, _ = fmt.Fprintf(os.Stderr, "[INFO] Dumping Path '%s':\n", apiPath)
return u.PrintRawAPIJSON(apiPath)
default:
@ -56,7 +56,7 @@ func (u *UnifiPoller) dumpSitesJSON(path, name string, sites unifi.Sites) error
// PrintRawAPIJSON prints the raw json for a user-provided path on a UniFi Controller.
func (u *UnifiPoller) PrintRawAPIJSON(apiPath string) error {
body, err := u.GetJSON(apiPath)
body, err := u.Unifi.GetJSON(apiPath)
fmt.Println(string(body))
return err
}

View File

@ -11,7 +11,7 @@ import (
func (u *UnifiPoller) LogError(err error, prefix string) {
if err != nil {
u.errorCount++
_ = log.Output(2, fmt.Sprintf("[ERROR] (%v/%v) %v: %v", u.errorCount, u.MaxErrors, prefix, err))
_ = log.Output(2, fmt.Sprintf("[ERROR] (%v/%v) %v: %v", u.errorCount, u.Config.MaxErrors, prefix, err))
}
}
@ -27,14 +27,14 @@ func StringInSlice(str string, slice []string) bool {
// Logf prints a log entry if quiet is false.
func (u *UnifiPoller) Logf(m string, v ...interface{}) {
if !u.Quiet {
if !u.Config.Quiet {
_ = log.Output(2, fmt.Sprintf("[INFO] "+m, v...))
}
}
// LogDebugf prints a debug log entry if debug is true and quite is false
func (u *UnifiPoller) LogDebugf(m string, v ...interface{}) {
if u.Debug && !u.Quiet {
if u.Config.Debug && !u.Config.Quiet {
_ = log.Output(2, fmt.Sprintf("[DEBUG] "+m, v...))
}
}

View File

@ -21,8 +21,9 @@ import (
// Parses flags, parses config and executes Run().
func Start() error {
log.SetFlags(log.LstdFlags)
up := &UnifiPoller{}
if up.ParseFlags(os.Args[1:]); up.ShowVer {
up := &UnifiPoller{Flag: &Flag{}}
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.
}
@ -33,18 +34,18 @@ func Start() error {
return up.Run()
}
// ParseFlags runs the parser.
func (u *UnifiPoller) ParseFlags(args []string) {
u.Flag = pflag.NewFlagSet("unifi-poller", pflag.ExitOnError)
u.Flag.Usage = func() {
// Parse turns CLI arguments into data structures. Called by Start() on startup.
func (f *Flag) Parse(args []string) {
f.FlagSet = pflag.NewFlagSet("unifi-poller", pflag.ExitOnError)
f.Usage = func() {
fmt.Println("Usage: unifi-poller [--config=filepath] [--version]")
u.Flag.PrintDefaults()
f.PrintDefaults()
}
u.Flag.StringVarP(&u.DumpJSON, "dumpjson", "j", "",
f.StringVarP(&f.DumpJSON, "dumpjson", "j", "",
"This debug option prints a json payload and exits. See man page for more.")
u.Flag.StringVarP(&u.ConfigFile, "config", "c", DefaultConfFile, "Poller Config File (TOML Format)")
u.Flag.BoolVarP(&u.ShowVer, "version", "v", false, "Print the version and exit")
_ = u.Flag.Parse(args)
f.StringVarP(&f.ConfigFile, "config", "c", DefaultConfFile, "Poller Config File (TOML Format)")
f.BoolVarP(&f.ShowVer, "version", "v", false, "Print the version and exit")
_ = f.FlagSet.Parse(args)
}
// GetConfig parses and returns our configuration data.
@ -60,17 +61,17 @@ func (u *UnifiPoller) GetConfig() error {
UnifiBase: defaultUnifURL,
Interval: Duration{defaultInterval},
Sites: []string{"default"},
Quiet: u.DumpJSON != "",
Quiet: u.Flag.DumpJSON != "",
}
u.Logf("Loading Configuration File: %s", u.ConfigFile)
switch buf, err := ioutil.ReadFile(u.ConfigFile); {
u.Logf("Loading Configuration File: %s", u.Flag.ConfigFile)
switch buf, err := ioutil.ReadFile(u.Flag.ConfigFile); {
case err != nil:
return err
case strings.Contains(u.ConfigFile, ".json"):
case strings.Contains(u.Flag.ConfigFile, ".json"):
return json.Unmarshal(buf, u.Config)
case strings.Contains(u.ConfigFile, ".xml"):
case strings.Contains(u.Flag.ConfigFile, ".xml"):
return xml.Unmarshal(buf, u.Config)
case strings.Contains(u.ConfigFile, ".yaml"):
case strings.Contains(u.Flag.ConfigFile, ".yaml"):
return yaml.Unmarshal(buf, u.Config)
default:
return toml.Unmarshal(buf, u.Config)
@ -79,10 +80,10 @@ func (u *UnifiPoller) GetConfig() error {
// Run invokes all the application logic and routines.
func (u *UnifiPoller) Run() (err error) {
if u.DumpJSON != "" {
if u.Flag.DumpJSON != "" {
return u.DumpJSONPayload()
}
if u.Debug {
if u.Config.Debug {
log.SetFlags(log.Lshortfile | log.Lmicroseconds | log.Ldate)
u.LogDebugf("Debug Logging Enabled")
}
@ -90,12 +91,13 @@ func (u *UnifiPoller) Run() (err error) {
if err = u.GetUnifi(); err != nil {
return err
}
u.Logf("Polling UniFi Controller at %s v%s as user %s. Sites: %v", u.UnifiBase, u.ServerVersion, u.UnifiUser, u.Sites)
u.Logf("Polling UniFi Controller at %s v%s as user %s. Sites: %v",
u.Config.UnifiBase, u.Unifi.ServerVersion, u.Config.UnifiUser, u.Config.Sites)
if err = u.GetInfluxDB(); err != nil {
return err
}
u.Logf("Logging Measurements to InfluxDB at %s as user %s", u.InfluxURL, u.InfluxUser)
switch strings.ToLower(u.Mode) {
u.Logf("Logging Measurements to InfluxDB at %s as user %s", u.Config.InfluxURL, u.Config.InfluxUser)
switch strings.ToLower(u.Config.Mode) {
case "influxlambda", "lambdainflux", "lambda_influx", "influx_lambda":
u.LogDebugf("Lambda Mode Enabled")
u.LastCheck = time.Now()
@ -107,10 +109,10 @@ func (u *UnifiPoller) Run() (err error) {
// GetInfluxDB returns an InfluxDB interface.
func (u *UnifiPoller) GetInfluxDB() (err error) {
u.Client, err = influx.NewHTTPClient(influx.HTTPConfig{
Addr: u.InfluxURL,
Username: u.InfluxUser,
Password: u.InfluxPass,
u.Influx, err = influx.NewHTTPClient(influx.HTTPConfig{
Addr: u.Config.InfluxURL,
Username: u.Config.InfluxUser,
Password: u.Config.InfluxPass,
})
if err != nil {
return fmt.Errorf("influxdb: %v", err)
@ -122,10 +124,10 @@ func (u *UnifiPoller) GetInfluxDB() (err error) {
func (u *UnifiPoller) GetUnifi() (err error) {
// Create an authenticated session to the Unifi Controller.
u.Unifi, err = unifi.NewUnifi(&unifi.Config{
User: u.UnifiUser,
Pass: u.UnifiPass,
URL: u.UnifiBase,
VerifySSL: u.VerifySSL,
User: u.Config.UnifiUser,
Pass: u.Config.UnifiPass,
URL: u.Config.UnifiBase,
VerifySSL: u.Config.VerifySSL,
ErrorLog: u.LogErrorf, // Log all errors.
DebugLog: u.LogDebugf, // Log debug messages.
})

View File

@ -13,10 +13,10 @@ import (
// CheckSites makes sure the list of provided sites exists on the controller.
// This does not run in Lambda (run-once) mode.
func (u *UnifiPoller) CheckSites() error {
if strings.Contains(strings.ToLower(u.Mode), "lambda") {
if strings.Contains(strings.ToLower(u.Config.Mode), "lambda") {
return nil // Skip this in lambda mode.
}
sites, err := u.GetSites()
sites, err := u.Unifi.GetSites()
if err != nil {
return err
}
@ -25,12 +25,12 @@ func (u *UnifiPoller) CheckSites() error {
msg = append(msg, site.Name+" ("+site.Desc+")")
}
u.Logf("Found %d site(s) on controller: %v", len(msg), strings.Join(msg, ", "))
if StringInSlice("all", u.Sites) {
u.Sites = []string{"all"}
if StringInSlice("all", u.Config.Sites) {
u.Config.Sites = []string{"all"}
return nil
}
FIRST:
for _, s := range u.Sites {
for _, s := range u.Config.Sites {
for _, site := range sites {
if s == site.Name {
continue FIRST
@ -45,14 +45,15 @@ FIRST:
// PollController runs forever, polling UniFi, and pushing to influx.
// This is started by Run() after everything checks out.
func (u *UnifiPoller) PollController() error {
log.Println("[INFO] Everything checks out! Poller started, interval:", u.Interval.Round(time.Second))
ticker := time.NewTicker(u.Interval.Round(time.Second))
interval := u.Config.Interval.Round(time.Second)
log.Println("[INFO] Everything checks out! Poller started, interval:", interval)
ticker := time.NewTicker(interval)
for u.LastCheck = range ticker.C {
var err error
if u.ReAuth {
if u.Config.ReAuth {
u.LogDebugf("Re-authenticating to UniFi Controller")
// Some users need to re-auth every interval because the cookie times out.
if err = u.Login(); err != nil {
if err = u.Unifi.Login(); err != nil {
u.LogError(err, "re-authenticating")
}
}
@ -60,8 +61,8 @@ func (u *UnifiPoller) PollController() error {
// Only run this if the authentication procedure didn't return error.
_ = u.CollectAndReport()
}
if u.MaxErrors >= 0 && u.errorCount > u.MaxErrors {
return fmt.Errorf("reached maximum error count, stopping poller (%d > %d)", u.errorCount, u.MaxErrors)
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 nil
@ -93,18 +94,18 @@ func (u *UnifiPoller) CollectMetrics() (*Metrics, error) {
// Get the sites we care about.
m.Sites, err = u.GetFilteredSites()
u.LogError(err, "unifi.GetSites()")
if u.CollectIDS {
if u.Config.CollectIDS {
// Check back in time since twice the interval. Dups are discarded by InfluxDB.
m.IDSList, err = u.GetIDS(m.Sites, time.Now().Add(2*u.Interval.Duration), time.Now())
m.IDSList, err = u.Unifi.GetIDS(m.Sites, time.Now().Add(2*u.Config.Interval.Duration), time.Now())
u.LogError(err, "unifi.GetIDS()")
}
// Get all the points.
m.Clients, err = u.GetClients(m.Sites)
m.Clients, err = u.Unifi.GetClients(m.Sites)
u.LogError(err, "unifi.GetClients()")
m.Devices, err = u.GetDevices(m.Sites)
m.Devices, err = u.Unifi.GetDevices(m.Sites)
u.LogError(err, "unifi.GetDevices()")
// Make a new Influx Points Batcher.
m.BatchPoints, err = influx.NewBatchPoints(influx.BatchPointsConfig{Database: u.InfluxDB})
m.BatchPoints, err = influx.NewBatchPoints(influx.BatchPointsConfig{Database: u.Config.InfluxDB})
u.LogError(err, "influx.NewBatchPoints")
return m, err
}
@ -139,7 +140,7 @@ func (u *UnifiPoller) ReportMetrics(metrics *Metrics) error {
for _, err := range metrics.ProcessPoints() {
u.LogError(err, "asset.Points()")
}
err := u.Write(metrics.BatchPoints)
err := u.Influx.Write(metrics.BatchPoints)
if err != nil {
return fmt.Errorf("influxdb.Write(points): %v", err)
}
@ -150,7 +151,7 @@ func (u *UnifiPoller) ReportMetrics(metrics *Metrics) error {
fields += len(i)
}
idsMsg := ""
if u.CollectIDS {
if u.Config.CollectIDS {
idsMsg = fmt.Sprintf("IDS Events: %d, ", len(metrics.IDSList))
}
u.Logf("UniFi Measurements Recorded. Sites: %d, Clients: %d, "+
@ -213,16 +214,16 @@ func (m *Metrics) ProcessPoints() []error {
// Omits requested but unconfigured sites. Grabs the full list from the
// controller and returns the sites provided in the config file.
func (u *UnifiPoller) GetFilteredSites() (unifi.Sites, error) {
sites, err := u.GetSites()
sites, err := u.Unifi.GetSites()
if err != nil {
return nil, err
} else if len(u.Sites) < 1 || StringInSlice("all", u.Sites) {
} else if len(u.Config.Sites) < 1 || StringInSlice("all", u.Config.Sites) {
return sites, nil
}
var i int
for _, s := range sites {
// Only include valid sites in the request filter.
if StringInSlice(s.Name, u.Sites) {
if StringInSlice(s.Name, u.Config.Sites) {
sites[i] = s
i++
}