Add support for multiple sites.

This commit is contained in:
David Newhall II 2019-04-15 01:26:24 -07:00
parent 4bebe6fbe8
commit 32dfaa2508
5 changed files with 60 additions and 18 deletions

View File

@ -2,12 +2,12 @@
[[projects]] [[projects]]
digest = "1:e7f0acf99afe9a7b03d270164bd2976b687e1aef02ab3a0abd8db0b4de44b817" digest = "1:9543804b8019077cb769bc3b15bfb0640fc732f8b061df323a973171b979d02f"
name = "github.com/golift/unifi" name = "github.com/golift/unifi"
packages = ["."] packages = ["."]
pruneopts = "UT" pruneopts = "UT"
revision = "f8fec42fbe169dceb69f15276a2323fb007a4539" revision = "1779d2c07abd748d0d23868356ef517ac2e00e7e"
version = "v1.0.0" version = "v1.0.1"
[[projects]] [[projects]]
branch = "master" branch = "master"

View File

@ -30,6 +30,10 @@ unifi-poller(1) -- Utility to poll Unifi Metrics and drop them into InfluxDB
`Config File Parameters` `Config File Parameters`
`sites` default: ["default"]
This list of strings should represent the names of sites on the unifi
controller that should be polled for data.
`interval` default: 30s `interval` default: 30s
How often to poll the controller for updated client and device data. How often to poll the controller for updated client and device data.
The Unifi Controller only updates traffic stats about every 30 seconds. The Unifi Controller only updates traffic stats about every 30 seconds.

View File

@ -33,6 +33,7 @@ type Config struct {
UnifiUser string `json:"unifi_user" toml:"unifi_user" xml:"unifi_user" yaml:"unifi_user"` UnifiUser string `json:"unifi_user" toml:"unifi_user" xml:"unifi_user" yaml:"unifi_user"`
UnifiPass string `json:"unifi_pass" toml:"unifi_pass" xml:"unifi_pass" yaml:"unifi_pass"` UnifiPass string `json:"unifi_pass" toml:"unifi_pass" xml:"unifi_pass" yaml:"unifi_pass"`
UnifiBase string `json:"unifi_url" toml:"unifi_url" xml:"unifi_url" yaml:"unifi_url"` UnifiBase string `json:"unifi_url" toml:"unifi_url" xml:"unifi_url" yaml:"unifi_url"`
Sites []string `json:"sites" toml:"sites" xml:"sites" yaml:"sites"`
} }
// Dur is used to UnmarshalTOML into a time.Duration value. // Dur is used to UnmarshalTOML into a time.Duration value.

View File

@ -5,6 +5,7 @@ import (
"io/ioutil" "io/ioutil"
"log" "log"
"os" "os"
"strings"
"time" "time"
"github.com/golift/unifi" "github.com/golift/unifi"
@ -86,6 +87,7 @@ func GetConfig(configFile string) (Config, error) {
Debug: defaultDebug, Debug: defaultDebug,
Quiet: defaultQuiet, Quiet: defaultQuiet,
Interval: Dur{value: defaultInterval}, Interval: Dur{value: defaultInterval},
Sites: []string{"default"},
} }
if buf, err := ioutil.ReadFile(configFile); err != nil { if buf, err := ioutil.ReadFile(configFile); err != nil {
return config, err return config, err
@ -102,9 +104,11 @@ func (c *Config) PollUnifiController(controller *unifi.Unifi, infdb influx.Clien
log.Println("[INFO] Everyting checks out! Beginning Poller Routine.") log.Println("[INFO] Everyting checks out! Beginning Poller Routine.")
ticker := time.NewTicker(c.Interval.value) ticker := time.NewTicker(c.Interval.value)
for range ticker.C { for range ticker.C {
if clients, err := controller.GetClients(); err != nil { if sites, err := filterSites(controller, c.Sites); err != nil {
logErrors([]error{err}, "uni.GetSites()")
} else if clients, err := controller.GetClients(sites); err != nil {
logErrors([]error{err}, "uni.GetClients()") logErrors([]error{err}, "uni.GetClients()")
} else if devices, err := controller.GetDevices(); err != nil { } else if devices, err := controller.GetDevices(sites); err != nil {
logErrors([]error{err}, "uni.GetDevices()") logErrors([]error{err}, "uni.GetDevices()")
} else if bp, err := influx.NewBatchPoints(influx.BatchPointsConfig{Database: c.InfluxDB}); err != nil { } else if bp, err := influx.NewBatchPoints(influx.BatchPointsConfig{Database: c.InfluxDB}); err != nil {
logErrors([]error{err}, "influx.NewBatchPoints") logErrors([]error{err}, "influx.NewBatchPoints")
@ -113,12 +117,31 @@ func (c *Config) PollUnifiController(controller *unifi.Unifi, infdb influx.Clien
} else if err := infdb.Write(bp); err != nil { } else if err := infdb.Write(bp); err != nil {
logErrors([]error{err}, "infdb.Write(bp)") logErrors([]error{err}, "infdb.Write(bp)")
} else if !c.Quiet { } else if !c.Quiet {
log.Println("[INFO] Logged Unifi States. Clients:", len(clients.UCLs), "- Wireless APs:", log.Printf("[INFO] Logged Unifi States. Sites: %d Clients: %d, Wireless APs: %d, Gateways: %d, Switches: %d",
len(devices.UAPs), "Gateways:", len(devices.USGs), "Switches:", len(devices.USWs)) 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) ([]string, error) {
sites, err := controller.GetSites()
if err != nil {
return filter, err
} else if len(filter) < 1 || StringInSlice("all", filter) {
return sites, nil
}
output := []string{}
for _, s := range filter {
// Do not return requested sites that are not configured.
if StringInSlice(s, sites) {
output = append(output, s)
}
}
return output, nil
}
// batchPoints combines all device and client data into influxdb data points. // batchPoints combines all device and client data into influxdb data points.
func batchPoints(devices *unifi.Devices, clients *unifi.Clients, batchPoints influx.BatchPoints) (errs []error) { func batchPoints(devices *unifi.Devices, clients *unifi.Clients, batchPoints influx.BatchPoints) (errs []error) {
process := func(asset Asset) error { process := func(asset Asset) error {
@ -162,3 +185,13 @@ func logErrors(errs []error, prefix string) {
} }
} }
} }
// 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
}

View File

@ -3,6 +3,10 @@
# commented lines are defaults, uncomment to change. # # commented lines are defaults, uncomment to change. #
########################################################## ##########################################################
# If the controller has more than one site, specify which sites to poll here.
# If only one site, "default" is likely the correct name.
#sites = ["default"]
# The Unifi Controller only updates traffic stats about every 30 seconds. # The Unifi Controller only updates traffic stats about every 30 seconds.
# Setting this to something lower may lead to "zeros" in your data. You've been warned. # Setting this to something lower may lead to "zeros" in your data. You've been warned.
#interval = "30s" #interval = "30s"