Merge pull request #18 from davidnewhall/dn2_multi_site

Add support for multiple sites.

Gopkg.lock

@@ -2,12 +2,12 @@
 
 
 [[projects]]
-  digest = "1:e7f0acf99afe9a7b03d270164bd2976b687e1aef02ab3a0abd8db0b4de44b817"
+  digest = "1:dbd1c953926f2aca29aeb24f13af5875a3b31f09ed7093cb18a5409d4b9fc9f1"
   name = "github.com/golift/unifi"
   packages = ["."]
   pruneopts = "UT"
-  revision = "f8fec42fbe169dceb69f15276a2323fb007a4539"
+  revision = "034a62e1cb6ef13e724cfab9b8f87dbb130d493c"
-  version = "v1.0.0"
+  version = "v2.0.0"
 
 [[projects]]
   branch = "master"
@@ -19,7 +19,7 @@
     "v2",
   ]
   pruneopts = "UT"
-  revision = "16c852ea613fa2d42fcdccc9a8b0802a8bdc6140"
+  revision = "8ff2fc3824fcb533795f9a2f233275f0bb18d6c5"
 
 [[projects]]
   digest = "1:b56c589214f01a5601e0821387db484617392d0042f26234bf2da853a2f498a1"

cmd/unifi-poller/README.md

@@ -30,6 +30,11 @@ 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 will be polled for data. Pass `all` in the list to
+        poll all sites.
+
     `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.

cmd/unifi-poller/config.go

@@ -7,32 +7,30 @@ var Version = "v1.0.0"
 
 const (
 	// App defaults in case they're missing from the config.
-	defaultConfFile  = "/usr/local/etc/unifi-poller/up.conf"
+	defaultConfFile = "/usr/local/etc/unifi-poller/up.conf"
-	defaultInterval  = 30 * time.Second
+	defaultInterval = 30 * time.Second
-	defaultInfxDb    = "unifi"
+	defaultInfxDb   = "unifi"
-	defaultInfxUser  = "unifi"
+	defaultInfxUser = "unifi"
-	defaultInfxPass  = "unifi"
+	defaultInfxPass = "unifi"
-	defaultInfxURL   = "http://127.0.0.1:8086"
+	defaultInfxURL  = "http://127.0.0.1:8086"
-	defaultUnifUser  = "influx"
+	defaultUnifUser = "influx"
-	defaultUnifURL   = "https://127.0.0.1:8443"
+	defaultUnifURL  = "https://127.0.0.1:8443"
-	defaultVerifySSL = false
-	defaultDebug     = false
-	defaultQuiet     = false
 )
 
 // Config represents the data needed to poll a controller and report to influxdb.
 type Config struct {
-	Interval   Dur    `json:"interval" toml:"interval" xml:"interval" yaml:"interval"`
+	Interval   Dur      `json:"interval,_omitempty" toml:"interval,_omitempty" xml:"interval" yaml:"interval"`
-	Debug      bool   `json:"debug" toml:"debug" xml:"debug" yaml:"debug"`
+	Debug      bool     `json:"debug" toml:"debug" xml:"debug" yaml:"debug"`
-	Quiet      bool   `json:"quiet" toml:"quiet" xml:"quiet" yaml:"quiet"`
+	Quiet      bool     `json:"quiet" toml:"quiet" xml:"quiet" yaml:"quiet"`
-	VerifySSL  bool   `json:"verify_ssl" toml:"verify_ssl" xml:"verify_ssl" yaml:"verify_ssl"`
+	VerifySSL  bool     `json:"verify_ssl" toml:"verify_ssl" xml:"verify_ssl" yaml:"verify_ssl"`
-	InfluxURL  string `json:"influx_url" toml:"influx_url" xml:"influx_url" yaml:"influx_url"`
+	InfluxURL  string   `json:"influx_url,_omitempty" toml:"influx_url,_omitempty" xml:"influx_url" yaml:"influx_url"`
-	InfluxUser string `json:"influx_user" toml:"influx_user" xml:"influx_user" yaml:"influx_user"`
+	InfluxUser string   `json:"influx_user,_omitempty" toml:"influx_user,_omitempty" xml:"influx_user" yaml:"influx_user"`
-	InfluxPass string `json:"influx_pass" toml:"influx_pass" xml:"influx_pass" yaml:"influx_pass"`
+	InfluxPass string   `json:"influx_pass,_omitempty" toml:"influx_pass,_omitempty" xml:"influx_pass" yaml:"influx_pass"`
-	InfluxDB   string `json:"influx_db" toml:"influx_db" xml:"influx_db" yaml:"influx_db"`
+	InfluxDB   string   `json:"influx_db,_omitempty" toml:"influx_db,_omitempty" xml:"influx_db" yaml:"influx_db"`
-	UnifiUser  string `json:"unifi_user" toml:"unifi_user" xml:"unifi_user" yaml:"unifi_user"`
+	UnifiUser  string   `json:"unifi_user,_omitempty" toml:"unifi_user,_omitempty" 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,_omitempty" toml:"unifi_pass,_omitempty" 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,_omitempty" toml:"unifi_url,_omitempty" xml:"unifi_url" yaml:"unifi_url"`
+	Sites      []string `json:"sites,_omitempty" toml:"sites,_omitempty" xml:"sites" yaml:"sites"`
 }
 
 // Dur is used to UnmarshalTOML into a time.Duration value.

cmd/unifi-poller/main.go

@@ -5,6 +5,7 @@ import (
 	"io/ioutil"
 	"log"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/golift/unifi"
@@ -27,7 +28,7 @@ func main() {
 		log.Fatalf("Config Error '%v': %v", configFile, err)
 	}
 	// Create an authenticated session to the Unifi Controller.
-	controller, err := unifi.GetController(config.UnifiUser, config.UnifiPass, config.UnifiBase, config.VerifySSL)
+	controller, err := unifi.NewUnifi(config.UnifiUser, config.UnifiPass, config.UnifiBase, config.VerifySSL)
 	if err != nil {
 		log.Fatalln("Unifi Controller Error:", err)
 	} else if !config.Quiet {
@@ -52,7 +53,7 @@ func main() {
 		controller.DebugLog = nil
 	} else {
 		log.Println("Logging Unifi Metrics to InfluXDB @", config.InfluxURL, "as user", config.InfluxUser)
-		log.Println("Polling Unifi Controller, interval:", config.Interval.value)
+		log.Printf("Polling Unifi Controller (sites %v), interval: %v", config.Sites, config.Interval.value)
 	}
 	config.PollUnifiController(controller, infdb)
 }
@@ -82,10 +83,8 @@ func GetConfig(configFile string) (Config, error) {
 		UnifiUser:  defaultUnifUser,
 		UnifiPass:  os.Getenv("UNIFI_PASSWORD"),
 		UnifiBase:  defaultUnifURL,
-		VerifySSL:  defaultVerifySSL,
-		Debug:      defaultDebug,
-		Quiet:      defaultQuiet,
 		Interval:   Dur{value: defaultInterval},
+		Sites:      []string{"default"},
 	}
 	if buf, err := ioutil.ReadFile(configFile); err != nil {
 		return config, err
@@ -102,9 +101,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 +114,32 @@ 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:",
+			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) ([]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, batchPoints influx.BatchPoints) (errs []error) {
 	process := func(asset Asset) error {
@@ -162,3 +183,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
+}

grafana-dashboards/unifi-clients-grafana-dash.json

@@ -3,7 +3,7 @@
     "list": [
       {
         "builtIn": 1,
-        "datasource": "-- Grafana --",
+        "datasource": "Unifi",
         "enable": true,
         "hide": true,
         "iconColor": "rgba(0, 211, 255, 1)",

grafana-dashboards/unifi-uap-grafana-dash.json

@@ -57,7 +57,7 @@
     "list": [
       {
         "builtIn": 1,
-        "datasource": "-- Grafana --",
+        "datasource": "${DS_UNIFI}",
         "enable": true,
         "hide": true,
         "iconColor": "rgba(0, 211, 255, 1)",

grafana-dashboards/unifi-usg-grafana-dash.json

@@ -57,7 +57,7 @@
     "list": [
       {
         "builtIn": 1,
-        "datasource": "-- Grafana --",
+        "datasource": "${DS_UNIFI}",
         "enable": true,
         "hide": true,
         "iconColor": "rgba(0, 211, 255, 1)",

up.conf.example

@@ -3,6 +3,11 @@
 # 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.
+# Set this to ["all"] to poll all sites, no matter their names.
+#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"