Add support for multiple sites.

integrations/promunifi/Gopkg.lock

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

integrations/promunifi/cmd/unifi-poller/README.md

@@ -30,6 +30,10 @@ unifi-poller(1) -- Utility to poll Unifi Metrics and drop them into InfluxDB
 
 `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
         How often to poll the controller for updated client and device data.
         The Unifi Controller only updates traffic stats about every 30 seconds.

integrations/promunifi/cmd/unifi-poller/config.go

@@ -33,6 +33,7 @@ type Config struct {
 	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"`
 	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.

integrations/promunifi/cmd/unifi-poller/main.go

@@ -5,6 +5,7 @@ import (
 	"io/ioutil"
 	"log"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/golift/unifi"
@@ -86,6 +87,7 @@ func GetConfig(configFile string) (Config, error) {
 		Debug:      defaultDebug,
 		Quiet:      defaultQuiet,
 		Interval:   Dur{value: defaultInterval},
+		Sites:      []string{"default"},
 	}
 	if buf, err := ioutil.ReadFile(configFile); err != nil {
 		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.")
 	ticker := time.NewTicker(c.Interval.value)
 	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()")
-		} else if devices, err := controller.GetDevices(); err != nil {
+		} else if devices, err := controller.GetDevices(sites); 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")
@@ -113,12 +117,31 @@ func (c *Config) PollUnifiController(controller *unifi.Unifi, infdb influx.Clien
 		} else if err := infdb.Write(bp); err != nil {
 			logErrors([]error{err}, "infdb.Write(bp)")
 		} else if !c.Quiet {
-			log.Println("[INFO] Logged Unifi States. Clients:", len(clients.UCLs), "- Wireless APs:",
-				len(devices.UAPs), "Gateways:", len(devices.USGs), "Switches:", len(devices.USWs))
+			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) ([]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.
 func batchPoints(devices *unifi.Devices, clients *unifi.Clients, batchPoints influx.BatchPoints) (errs []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
+}

integrations/promunifi/up.conf.example

@@ -3,6 +3,10 @@
 # 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.
 # Setting this to something lower may lead to "zeros" in your data. You've been warned.
 #interval = "30s"