From bac02d938a6d2b4d7730e718e8e57b694c0b16fb Mon Sep 17 00:00:00 2001 From: unifi-poller-bot Date: Wed, 12 Jun 2019 21:02:40 -0700 Subject: [PATCH] Split things up a bit --- cmd/unifi-poller/config.go | 8 +- cmd/unifi-poller/helpers.go | 35 +++++ cmd/unifi-poller/main.go | 255 ++++++++---------------------------- cmd/unifi-poller/unifi.go | 132 +++++++++++++++++++ 4 files changed, 225 insertions(+), 205 deletions(-) create mode 100644 cmd/unifi-poller/helpers.go create mode 100644 cmd/unifi-poller/unifi.go diff --git a/cmd/unifi-poller/config.go b/cmd/unifi-poller/config.go index 203a88fd..089a8520 100644 --- a/cmd/unifi-poller/config.go +++ b/cmd/unifi-poller/config.go @@ -3,6 +3,7 @@ package main import ( "time" + influx "github.com/influxdata/influxdb1-client/v2" "github.com/spf13/pflag" ) @@ -21,13 +22,18 @@ const ( defaultUnifURL = "https://127.0.0.1:8443" ) +// Asset is used to give all devices and clients a common interface. +type Asset interface { + Points() ([]*influx.Point, error) +} + // UnifiPoller contains the application startup data. type UnifiPoller struct { ConfigFile string DumpJSON string ShowVer bool + Flag *pflag.FlagSet *Config - Flag *pflag.FlagSet } // Config represents the data needed to poll a controller and report to influxdb. diff --git a/cmd/unifi-poller/helpers.go b/cmd/unifi-poller/helpers.go new file mode 100644 index 00000000..df9fa5ca --- /dev/null +++ b/cmd/unifi-poller/helpers.go @@ -0,0 +1,35 @@ +package main + +import ( + "log" + "strings" +) + +// hasErr checks a list of errors for a non-nil. +func hasErr(errs []error) bool { + for _, err := range errs { + if err != nil { + return true + } + } + return false +} + +// logErrors writes a slice of errors, with a prefix, to log-out. +func logErrors(errs []error, prefix string) { + for _, err := range errs { + if err != nil { + log.Println("[ERROR]", prefix+":", err.Error()) + } + } +} + +// StringInSlice returns true if a string is in a slice. +func StringInSlice(str string, slc []string) bool { + for _, s := range slc { + if strings.EqualFold(s, str) { + return true + } + } + return false +} diff --git a/cmd/unifi-poller/main.go b/cmd/unifi-poller/main.go index 562dfc22..54f635f9 100644 --- a/cmd/unifi-poller/main.go +++ b/cmd/unifi-poller/main.go @@ -5,8 +5,6 @@ import ( "io/ioutil" "log" "os" - "strings" - "time" "github.com/golift/unifi" influx "github.com/influxdata/influxdb1-client/v2" @@ -15,11 +13,6 @@ import ( flag "github.com/spf13/pflag" ) -// Asset is used to give all devices and clients a common interface. -type Asset interface { - Points() ([]*influx.Point, error) -} - func main() { u := &UnifiPoller{} if u.ParseFlags(os.Args[1:]); u.ShowVer { @@ -30,19 +23,62 @@ func main() { u.Flag.Usage() log.Fatalf("[ERROR] config file '%v': %v", u.ConfigFile, err) } - if u.DumpJSON != "" { - if err := u.Config.DumpJSON(u.DumpJSON); err != nil { - log.Fatalln("[ERROR] dumping JSON:", err) - } - return - } - if err := u.Config.Run(); err != nil { + if err := u.Run(); err != nil { log.Fatalln("[ERROR]", err) } } +// ParseFlags runs the parser. +func (u *UnifiPoller) ParseFlags(args []string) { + u.Flag = flag.NewFlagSet("unifi-poller", flag.ExitOnError) + u.Flag.Usage = func() { + fmt.Println("Usage: unifi-poller [--config=filepath] [--version]") + u.Flag.PrintDefaults() + } + u.Flag.StringVarP(&u.DumpJSON, "dumpjson", "j", "", "This debug option prints the json payload for a device and exits.") + 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) +} + +// 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"), + UnifiBase: defaultUnifURL, + Interval: Dur{value: defaultInterval}, + Sites: []string{"default"}, + } + if buf, err := ioutil.ReadFile(u.ConfigFile); err != nil { + return err + // This is where the defaults in the config variable are overwritten. + } else if err := toml.Unmarshal(buf, u.Config); err != nil { + return err + } + if u.DumpJSON != "" { + u.Quiet = true + } + if !u.Config.Quiet { + log.Println("Loaded Configuration:", u.ConfigFile) + } + return nil +} + // Run invokes all the application logic and routines. -func (c *Config) Run() error { +func (u *UnifiPoller) Run() error { + c := u.Config + if u.DumpJSON != "" { + if err := c.DumpJSON(u.DumpJSON); err != nil { + log.Fatalln("[ERROR] dumping JSON:", err) + } + return nil + } log.Println("Unifi-Poller Starting Up! PID:", os.Getpid()) // Create an authenticated session to the Unifi Controller. controller, err := unifi.NewUnifi(c.UnifiUser, c.UnifiPass, c.UnifiBase, c.VerifySSL) @@ -79,192 +115,3 @@ func (c *Config) Run() error { c.PollUnifiController(controller, infdb) return nil } - -// ParseFlags runs the parser. -func (u *UnifiPoller) ParseFlags(args []string) { - u.Flag = flag.NewFlagSet("unifi-poller", flag.ExitOnError) - u.Flag.Usage = func() { - fmt.Println("Usage: unifi-poller [--config=filepath] [--version]") - u.Flag.PrintDefaults() - } - u.Flag.StringVarP(&u.DumpJSON, "dumpjson", "j", "", "This debug option prints the json payload for a device and exits.") - 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) -} - -// CheckSites makes sure the list of provided sites exists on the controller. -func (c *Config) CheckSites(controller *unifi.Unifi) error { - sites, err := controller.GetSites() - if err != nil { - return err - } - if !c.Quiet { - msg := []string{} - for _, site := range sites { - msg = append(msg, site.Name+" ("+site.Desc+")") - } - log.Printf("Found %d site(s) on controller: %v", len(msg), strings.Join(msg, ", ")) - } - if StringInSlice("all", c.Sites) { - return nil - } -FIRST: - for _, s := range c.Sites { - for _, site := range sites { - if s == site.Name { - continue FIRST - } - } - return errors.Errorf("configured site not found on controller: %v", s) - } - return nil -} - -// 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"), - UnifiBase: defaultUnifURL, - Interval: Dur{value: defaultInterval}, - Sites: []string{"default"}, - } - if buf, err := ioutil.ReadFile(u.ConfigFile); err != nil { - return err - // This is where the defaults in the config variable are overwritten. - } else if err := toml.Unmarshal(buf, u.Config); err != nil { - return err - } - if u.DumpJSON != "" { - u.Quiet = true - } - if !u.Config.Quiet { - log.Println("Loaded Configuration:", u.ConfigFile) - } - return nil -} - -// PollUnifiController runs forever, polling and pushing. -func (c *Config) PollUnifiController(controller *unifi.Unifi, infdb influx.Client) { - log.Println("[INFO] Everything checks out! Beginning Poller Routine.") - ticker := time.NewTicker(c.Interval.value) - - for range ticker.C { - sites, err := filterSites(controller, c.Sites) - if err != nil { - logErrors([]error{err}, "uni.GetSites()") - } - // Get all the points. - clients, err := controller.GetClients(sites) - if err != nil { - logErrors([]error{err}, "uni.GetClients()") - } - devices, err := controller.GetDevices(sites) - if err != nil { - logErrors([]error{err}, "uni.GetDevices()") - } - bp, err := influx.NewBatchPoints(influx.BatchPointsConfig{Database: c.InfluxDB}) - if err != nil { - logErrors([]error{err}, "influx.NewBatchPoints") - continue - } - // Batch all the points. - if errs := batchPoints(devices, clients, bp); errs != nil && hasErr(errs) { - logErrors(errs, "asset.Points()") - } - if err := infdb.Write(bp); err != nil { - logErrors([]error{err}, "infdb.Write(bp)") - } - if !c.Quiet { - log.Printf("[INFO] Logged Unifi States. Sites: %d Clients: %d, Wireless APs: %d, Gateways: %d, Switches: %d", - len(sites), len(clients.UCLs), len(devices.UAPs), len(devices.USGs), len(devices.USWs)) - } - } -} - -// filterSites returns a list of sites to fetch data for. -// Omits requested but unconfigured sites. -func filterSites(controller *unifi.Unifi, filter []string) ([]unifi.Site, error) { - sites, err := controller.GetSites() - if err != nil { - return nil, err - } else if len(filter) < 1 || StringInSlice("all", filter) { - return sites, nil - } - var i int - for _, s := range sites { - // Only include valid sites in the request filter. - if StringInSlice(s.Name, filter) { - sites[i] = s - i++ - } - } - return sites[:i], nil -} - -// batchPoints combines all device and client data into influxdb data points. -func batchPoints(devices *unifi.Devices, clients *unifi.Clients, bp influx.BatchPoints) (errs []error) { - process := func(asset Asset) error { - if asset == nil { - return nil - } - influxPoints, err := asset.Points() - if err != nil { - return err - } - bp.AddPoints(influxPoints) - return nil - } - if devices != 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)) - } - } - if clients != nil { - for _, asset := range clients.UCLs { - errs = append(errs, process(asset)) - } - } - return -} - -// hasErr checks a list of errors for a non-nil. -func hasErr(errs []error) bool { - for _, err := range errs { - if err != nil { - return true - } - } - return false -} - -// logErrors writes a slice of errors, with a prefix, to log-out. -func logErrors(errs []error, prefix string) { - for _, err := range errs { - if err != nil { - log.Println("[ERROR]", prefix+":", err.Error()) - } - } -} - -// StringInSlice returns true if a string is in a slice. -func StringInSlice(str string, slc []string) bool { - for _, s := range slc { - if strings.EqualFold(s, str) { - return true - } - } - return false -} diff --git a/cmd/unifi-poller/unifi.go b/cmd/unifi-poller/unifi.go new file mode 100644 index 00000000..3db402c6 --- /dev/null +++ b/cmd/unifi-poller/unifi.go @@ -0,0 +1,132 @@ +package main + +import ( + "log" + "strings" + "time" + + "github.com/golift/unifi" + influx "github.com/influxdata/influxdb1-client/v2" + "github.com/pkg/errors" +) + +// CheckSites makes sure the list of provided sites exists on the controller. +func (c *Config) CheckSites(controller *unifi.Unifi) error { + sites, err := controller.GetSites() + if err != nil { + return err + } + if !c.Quiet { + msg := []string{} + for _, site := range sites { + msg = append(msg, site.Name+" ("+site.Desc+")") + } + log.Printf("Found %d site(s) on controller: %v", len(msg), strings.Join(msg, ", ")) + } + if StringInSlice("all", c.Sites) { + return nil + } +FIRST: + for _, s := range c.Sites { + for _, site := range sites { + if s == site.Name { + continue FIRST + } + } + return errors.Errorf("configured site not found on controller: %v", s) + } + return nil +} + +// PollUnifiController runs forever, polling and pushing. +func (c *Config) PollUnifiController(controller *unifi.Unifi, infdb influx.Client) { + log.Println("[INFO] Everything checks out! Beginning Poller Routine.") + ticker := time.NewTicker(c.Interval.value) + + for range ticker.C { + // Get the sites we care about. + sites, err := filterSites(controller, c.Sites) + if err != nil { + logErrors([]error{err}, "uni.GetSites()") + } + // Get all the points. + clients, err := controller.GetClients(sites) + if err != nil { + logErrors([]error{err}, "uni.GetClients()") + } + devices, err := controller.GetDevices(sites) + if err != nil { + logErrors([]error{err}, "uni.GetDevices()") + } + // Make a new Points Batcher. + bp, err := influx.NewBatchPoints(influx.BatchPointsConfig{Database: c.InfluxDB}) + if err != nil { + logErrors([]error{err}, "influx.NewBatchPoints") + continue + } + // Batch (and send) all the points. + if errs := batchPoints(devices, clients, bp); errs != nil && hasErr(errs) { + logErrors(errs, "asset.Points()") + } + if err := infdb.Write(bp); err != nil { + logErrors([]error{err}, "infdb.Write(bp)") + } + // Talk about the data. + if !c.Quiet { + log.Printf("[INFO] Logged Unifi States. Sites: %d Clients: %d, Wireless APs: %d, Gateways: %d, Switches: %d", + len(sites), len(clients.UCLs), len(devices.UAPs), len(devices.USGs), len(devices.USWs)) + } + } +} + +// batchPoints combines all device and client data into influxdb data points. +func batchPoints(devices *unifi.Devices, clients *unifi.Clients, bp influx.BatchPoints) (errs []error) { + process := func(asset Asset) error { + if asset == nil { + return nil + } + influxPoints, err := asset.Points() + if err != nil { + return err + } + bp.AddPoints(influxPoints) + return nil + } + if devices != 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)) + } + } + if clients != nil { + for _, asset := range clients.UCLs { + errs = append(errs, process(asset)) + } + } + return +} + +// filterSites returns a list of sites to fetch data for. +// Omits requested but unconfigured sites. +func filterSites(controller *unifi.Unifi, filter []string) ([]unifi.Site, error) { + sites, err := controller.GetSites() + if err != nil { + return nil, err + } else if len(filter) < 1 || StringInSlice("all", filter) { + return sites, nil + } + var i int + for _, s := range sites { + // Only include valid sites in the request filter. + if StringInSlice(s.Name, filter) { + sites[i] = s + i++ + } + } + return sites[:i], nil +}