diff --git a/cmd/unifi-poller/main.go b/cmd/unifi-poller/main.go index a3182588..1742d3af 100644 --- a/cmd/unifi-poller/main.go +++ b/cmd/unifi-poller/main.go @@ -8,11 +8,16 @@ import ( "time" "github.com/golift/unifi" - influx "github.com/influxdata/influxdb/client/v2" + influx "github.com/influxdata/influxdb1-client/v2" "github.com/naoina/toml" flag "github.com/ogier/pflag" ) +// Asset is used to give all devices and clients a common interface. +type Asset interface { + Points() ([]*influx.Point, error) +} + func main() { configFile := parseFlags() log.Println("Unifi-Poller Starting Up! PID:", os.Getpid()) @@ -24,14 +29,14 @@ func main() { log.SetFlags(log.Lshortfile | log.Lmicroseconds | log.Ldate) log.Println("Debug Logging Enabled") } - log.Println("Loaded Configuration:", configFile) // Create an authenticated session to the Unifi Controller. - device, err := unifi.AuthController(config.UnifiUser, config.UnifiPass, config.UnifiBase, config.VerifySSL) + device, err := unifi.GetController(config.UnifiUser, config.UnifiPass, config.UnifiBase, config.VerifySSL) if err != nil { log.Fatalln("Unifi Controller Error:", err) } else if !config.Quiet { log.Println("Authenticated to Unifi Controller @", config.UnifiBase, "as user", config.UnifiUser) } + device.ErrorLog = log.Printf infdb, err := influx.NewHTTPClient(influx.HTTPConfig{ Addr: config.InfluxURL, Username: config.InfluxUser, @@ -41,13 +46,12 @@ func main() { log.Fatalln("InfluxDB Error:", err) } else if config.Quiet { // Doing it this way allows debug error logs (line numbers, etc) - unifi.Debug = false + device.DebugLog = log.Printf } else { log.Println("Logging Unifi Metrics to InfluXDB @", config.InfluxURL, "as user", config.InfluxUser) log.Println("Polling Unifi Controller, interval:", config.Interval.value) } - log.Println("Everyting checks out! Beginning Poller Routine.") - config.PollUnifiController(infdb, device, config.Quiet) + config.PollUnifiController(infdb, device) } func parseFlags() string { @@ -86,39 +90,65 @@ func GetConfig(configFile string) (Config, error) { } else if err := toml.Unmarshal(buf, &config); err != nil { return config, err } + log.Println("Loaded Configuration:", configFile) return config, nil } // PollUnifiController runs forever, polling and pushing. -func (c *Config) PollUnifiController(infdb influx.Client, device *unifi.AuthedReq, quiet bool) { +func (c *Config) PollUnifiController(infdb influx.Client, uni *unifi.Unifi) { + log.Println("Everyting checks out! Beginning Poller Routine.") ticker := time.NewTicker(c.Interval.value) for range ticker.C { - var clients, devices []unifi.Asset - var bp influx.BatchPoints - var err error - if clients, err = device.GetUnifiClientAssets(); err != nil { - log.Println("ERROR unifi.GetUnifiClientsAssets():", err) - } else if devices, err = device.GetUnifiDeviceAssets(); err != nil { - log.Println("ERROR unifi.GetUnifiDeviceAssets():", err) - } else if bp, err = influx.NewBatchPoints(influx.BatchPointsConfig{Database: c.InfluxDB}); err != nil { - log.Println("ERROR influx.NewBatchPoints:", err) - } - if err != nil { - continue - } - for _, asset := range append(clients, devices...) { - if pt, errr := asset.Points(); errr != nil { - log.Println("ERROR asset.Points():", errr) - } else { - bp.AddPoints(pt) - } - } - if err = infdb.Write(bp); err != nil { - log.Println("ERROR infdb.Write(bp):", err) - continue - } - if !quiet { - log.Println("Logged Unifi States. Clients:", len(clients), "- Devices:", len(devices)) + if clients, err := uni.GetClients(); err != nil { + logErrors([]error{err}, "uni.GetClients()") + } else if devices, err := uni.GetDevices(); err != nil { + logErrors([]error{err}, "uni.GetDevices()") + } else if bp, err := influx.NewBatchPoints(influx.BatchPointsConfig{Database: c.InfluxDB}); err != nil { + logErrors([]error{err}, "influx.NewBatchPoints") + } else if errs := batchPoints(devices, clients, bp); errs != nil { + logErrors(errs, "asset.Points()") + } else if err := infdb.Write(bp); err != nil { + logErrors([]error{err}, "infdb.Write(bp)") + } else if !c.Quiet { + log.Println("Logged Unifi States. Clients:", len(clients.UCLs), "- Wireless APs:", + len(devices.UAPs), "Gateways:", len(devices.USGs), "Switches:", len(devices.USWs)) + } + } +} + +// batchPoints combines all device and client data into influxdb data points. +func batchPoints(devices *unifi.Devices, clients *unifi.Clients, batchPoints influx.BatchPoints) (errs []error) { + process := func(asset Asset) error { + influxPoints, err := asset.Points() + if err != nil { + return err + } + batchPoints.AddPoints(influxPoints) + return nil + } + for _, asset := range devices.UAPs { + errs = append(errs, process(asset)) + } + for _, asset := range devices.USGs { + errs = append(errs, process(asset)) + } + for _, asset := range devices.USWs { + errs = append(errs, process(asset)) + } + for _, asset := range clients.UCLs { + errs = append(errs, process(asset)) + } + return +} + +// logErrors writes a slice of errors, with a prefix, to log-out. +func logErrors(errs []error, prefix string) { + if errs == nil { + return + } + for _, err := range errs { + if err != nil { + log.Println("[ERROR]", prefix+":", err.Error()) } } }