From 2a355c0677876029b219ccef1c69f8b4f61d1a8f Mon Sep 17 00:00:00 2001 From: David Newhall II Date: Sun, 25 Aug 2019 19:01:29 -0700 Subject: [PATCH] Remove struct inheritance. --- core/poller/unifipoller/config.go | 17 +++++--- core/poller/unifipoller/dumper.go | 22 +++++------ core/poller/unifipoller/helpers.go | 6 +-- core/poller/unifipoller/start.go | 62 +++++++++++++++--------------- core/poller/unifipoller/unifi.go | 43 +++++++++++---------- 5 files changed, 79 insertions(+), 71 deletions(-) diff --git a/core/poller/unifipoller/config.go b/core/poller/unifipoller/config.go index dad19c3d..39eac38e 100644 --- a/core/poller/unifipoller/config.go +++ b/core/poller/unifipoller/config.go @@ -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. diff --git a/core/poller/unifipoller/dumper.go b/core/poller/unifipoller/dumper.go index bb0fb0d8..ae110ceb 100644 --- a/core/poller/unifipoller/dumper.go +++ b/core/poller/unifipoller/dumper.go @@ -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 } diff --git a/core/poller/unifipoller/helpers.go b/core/poller/unifipoller/helpers.go index a138a052..c2e0fce4 100644 --- a/core/poller/unifipoller/helpers.go +++ b/core/poller/unifipoller/helpers.go @@ -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...)) } } diff --git a/core/poller/unifipoller/start.go b/core/poller/unifipoller/start.go index c6d360cb..145f1521 100644 --- a/core/poller/unifipoller/start.go +++ b/core/poller/unifipoller/start.go @@ -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. }) diff --git a/core/poller/unifipoller/unifi.go b/core/poller/unifipoller/unifi.go index ce674aa3..3cf0171f 100644 --- a/core/poller/unifipoller/unifi.go +++ b/core/poller/unifipoller/unifi.go @@ -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++ }