Restructure the code.

core/poller/.env.example

@@ -1,5 +1,6 @@
 UNIFI_ADDR="107.170.232.179"
 UNIFI_PORT="8443"
+# Go to Settings -> Admins and add (or use) a read-only user for this.
 UNIFI_USERNAME="username"
 UNIFI_PASSWORD="password"
 

core/poller/main.go

@@ -5,25 +5,48 @@ import (
 	"crypto/tls"
 	"encoding/json"
 	"errors"
+	"io"
 	"io/ioutil"
 	"log"
 	"net/http"
 	"net/http/cookiejar"
 	"os"
+	"strconv"
 	"time"
 
 	influx "github.com/influxdata/influxdb/client/v2"
 )
 
-// Response marshalls the payload from the controller.
+const (
-type Response struct {
+	// LoginPath is Unifi Controller Login API Path
-	Data []Client
+	LoginPath = "/api/login"
+	// ClientPath is Unifi Clients API Path
+	ClientPath = "/api/s/default/stat/sta"
+)
+
+// ClientResponse marshalls the payload from the controller.
+type ClientResponse struct {
+	Clients []Client `json:"data"`
 	Meta    struct {
-		Rc string
+		Rc string `json:"rc"`
-	}
+	} `json:"meta"`
+}
+
+// Config represents the data needed to poll a controller and report to influxdb.
+type Config struct {
+	Interval   time.Duration `json:"interval",toml:"interval",yaml:"interval"`
+	InfluxAddr string        `json:"influx_addr",toml:"influx_addr",yaml:"influx_addr"`
+	InfluxUser string        `json:"influx_user",toml:"influx_user",yaml:"influx_user"`
+	InfluxPass string        `json:"influx_pass",toml:"influx_pass",yaml:"influx_pass"`
+	InfluxDB   string        `json:"influx_db",toml:"influx_db",yaml:"influx_db"`
+	UnifiUser  string        `json:"unifi_user",toml:"unifi_user",yaml:"unifi_user"`
+	UnifiPass  string        `json:"unifi_pass"toml:"unifi_pass",yaml:"unifi_pass"`
+	UnifiBase  string        `json:"unifi_url",toml:"unifi_url",yaml:"unifi_url"`
+	uniClient  *http.Client
 }
 
 // DpiStat is for deep packet inspection stats.
+// Does not seem to exist in Unifi 5.7.20.
 type DpiStat struct {
 	App       int64
 	Cat       int64
@@ -38,91 +61,133 @@ type Client struct {
 	ID                  string    `json:"_id"`
 	IsGuestByUAP        bool      `json:"_is_guest_by_uap"`
 	IsGuestByUGW        bool      `json:"_is_guest_by_ugw"`
+	IsGuestByUSW        bool      `json:"_is_guest_by_usw"`
 	LastSeenByUAP       int64     `json:"_last_seen_by_uap"`
 	LastSeenByUGW       int64     `json:"_last_seen_by_ugw"`
+	LastSeenByUSW       int64     `json:"_last_seen_by_usw"`
 	UptimeByUAP         int64     `json:"_uptime_by_uap"`
 	UptimeByUGW         int64     `json:"_uptime_by_ugw"`
+	UptimeByUSW         int64     `json:"_uptime_by_usw"`
 	ApMac               string    `json:"ap_mac"`
 	AssocTime           int64     `json:"assoc_time"`
-	Authorized          bool
+	Authorized          bool      `json:"authorized"`
-	Bssid               string
+	Bssid               string    `json:"bssid"`
 	BytesR              int64     `json:"bytes-r"`
-	Ccq                 int64
+	Ccq                 int64     `json:"ccq"`
-	Channel             int64
+	Channel             int       `json:"channel"`
 	DpiStats            []DpiStat `json:"dpi_stats"`
 	DpiStatsLastUpdated int64     `json:"dpi_stats_last_updated"`
-	Essid               string
+	Essid               string    `json:"essid"`
 	FirstSeen           int64     `json:"first_seen"`
 	FixedIP             string    `json:"fixed_ip"`
 	Hostname            string    `json:"hostname"`
 	GwMac               string    `json:"gw_mac"`
 	IdleTime            int64     `json:"idle_time"`
-	IP                  string
+	IP                  string    `json:"ip"`
+	Is11R               bool      `json:"is_11r"`
 	IsGuest             bool      `json:"is_guest"`
 	IsWired             bool      `json:"is_wired"`
 	LastSeen            int64     `json:"last_seen"`
 	LatestAssocTime     int64     `json:"latest_assoc_time"`
-	Mac                 string
+	Mac                 string    `json:"mac"`
-	Name                string
+	Name                string    `json:"name"`
-	Network             string
+	Network             string    `json:"network"`
 	NetworkID           string    `json:"network_id"`
-	Noise               int64
+	Noise               int64     `json:"noise"`
-	Oui                 string
+	Note                string    `json:"note"`
+	Noted               bool      `json:"noted"`
+	Oui                 string    `json:"oui"`
 	PowersaveEnabled    bool      `json:"powersave_enabled"`
 	QosPolicyApplied    bool      `json:"qos_policy_applied"`
-	Radio               string
+	Radio               string    `json:"radio"`
 	RadioName           string    `json:"radio_name"`
 	RadioProto          string    `json:"radio_proto"`
 	RoamCount           int64     `json:"roam_count"`
-	Rssi                int64
+	Rssi                int64     `json:"rssi"`
 	RxBytes             int64     `json:"rx_bytes"`
 	RxBytesR            int64     `json:"rx_bytes-r"`
 	RxPackets           int64     `json:"rx_packets"`
 	RxRate              int64     `json:"rx_rate"`
-	Signal              int64
+	Signal              int64     `json:"signal"`
 	SiteID              string    `json:"site_id"`
+	SwDepth             int       `json:"sw_depth"`
+	SwMac               string    `json:"sw_mac"`
+	SwPort              int       `json:"sw_port"`
 	TxBytes             int64     `json:"tx_bytes"`
 	TxBytesR            int64     `json:"tx_bytes-r"`
 	TxPackets           int64     `json:"tx_packets"`
 	TxPower             int64     `json:"tx_power"`
 	TxRate              int64     `json:"tx_rate"`
-	Uptime              int64
+	Uptime              int64     `json:"uptime"`
 	UserID              string    `json:"user_id"`
-	Vlan                int64
+	UserGroupID         string    `json:"usergroup_id"`
+	UseFixedIP          bool      `json:"use_fixedip"`
+	Vlan                int       `json:"vlan"`
 }
 
-// Point generates a datapoint for InfluxDB.
+// Point generates a client's datapoint for InfluxDB.
 func (c Client) Point() *influx.Point {
+	if c.Name == "" && c.Hostname != "" {
+		c.Name = c.Hostname
+	} else if c.Hostname == "" && c.Name != "" {
+		c.Hostname = c.Name
+	} else if c.Hostname == "" && c.Name == "" {
+		c.Hostname = "-no-name-"
+		c.Name = "-no-name-"
+	}
 	tags := map[string]string{
+		"id":                 c.ID,
 		"mac":                c.Mac,
 		"user_id":            c.UserID,
 		"site_id":            c.SiteID,
 		"ip":                 c.IP,
+		"fixed_ip":           c.FixedIP,
 		"essid":              c.Essid,
+		"bssid":              c.Bssid,
 		"network":            c.Network,
+		"network_id":         c.NetworkID,
+		"usergroup_id":       c.UserGroupID,
 		"ap_mac":             c.ApMac,
+		"gw_mac":             c.GwMac,
+		"sw_mac":             c.SwMac,
+		"sw_port":            strconv.Itoa(c.SwPort),
+		"oui":                c.Oui,
 		"name":               c.Name,
 		"hostname":           c.Hostname,
 		"radio_name":         c.RadioName,
+		"radio":              c.Radio,
+		"radio_proto":        c.RadioProto,
+		"authorized":         strconv.FormatBool(c.Authorized),
+		"is_11r":             strconv.FormatBool(c.Is11R),
+		"is_wired":           strconv.FormatBool(c.IsWired),
+		"is_guest":           strconv.FormatBool(c.IsGuest),
+		"is_guest_by_uap":    strconv.FormatBool(c.IsGuestByUAP),
+		"is_guest_by_ugw":    strconv.FormatBool(c.IsGuestByUGW),
+		"is_guest_by_usw":    strconv.FormatBool(c.IsGuestByUSW),
+		"noted":              strconv.FormatBool(c.Noted),
+		"powersave_enabled":  strconv.FormatBool(c.PowersaveEnabled),
+		"qos_policy_applied": strconv.FormatBool(c.QosPolicyApplied),
+		"use_fixedip":        strconv.FormatBool(c.UseFixedIP),
+		"channel":            strconv.Itoa(c.Channel),
+		"vlan":               strconv.Itoa(c.Vlan),
 	}
 	fields := map[string]interface{}{
-		"is_guest_by_uap":        c.IsGuestByUAP,
+		"dpi_stats_last_updated": c.DpiStatsLastUpdated,
-		"is_guest_by_ugw":        c.IsGuestByUGW,
-		"authorized":             c.Authorized,
 		"last_seen_by_uap":       c.LastSeenByUAP,
 		"last_seen_by_ugw":       c.LastSeenByUGW,
+		"last_seen_by_usw":       c.LastSeenByUSW,
 		"uptime_by_uap":          c.UptimeByUAP,
 		"uptime_by_ugw":          c.UptimeByUGW,
+		"uptime_by_usw":          c.UptimeByUSW,
 		"assoc_time":             c.AssocTime,
 		"bytes_r":                c.BytesR,
 		"ccq":                    c.Ccq,
-		"channel":                c.Channel,
-		"dpi_stats_last_updated": c.DpiStatsLastUpdated,
 		"first_seen":             c.FirstSeen,
 		"idle_time":              c.IdleTime,
 		"last_seen":              c.LastSeen,
 		"latest_assoc_time":      c.LatestAssocTime,
 		"noise":                  c.Noise,
+		"note":                   c.Note,
 		"roam_count":             c.RoamCount,
 		"rssi":                   c.Rssi,
 		"rx_bytes":               c.RxBytes,
@@ -136,7 +201,6 @@ func (c Client) Point() *influx.Point {
 		"tx_power":               c.TxPower,
 		"tx_rate":                c.TxRate,
 		"uptime":                 c.Uptime,
-		"vlan":                   c.Vlan,
 	}
 
 	pt, err := influx.NewPoint("client_state", tags, fields, time.Now())
@@ -149,108 +213,116 @@ func (c Client) Point() *influx.Point {
 }
 
 func main() {
-	tickRate := os.Getenv("TICK_RATE")
+	config := GetConfig()
-	interval, err := time.ParseDuration(tickRate)
+	if err := config.AuthController(); err != nil {
-	if err != nil {
+		log.Fatal(err)
-		panic(err)
-	}
-
-	unifi, err := login()
-	if err != nil {
-		panic(err)
 	}
 	log.Println("Successfully authenticated to Unifi Controller!")
-	database := os.Getenv("INFLUXDB_DATABASE")
+
-	stats, err := influx.NewHTTPClient(influx.HTTPConfig{
+	infdb, err := influx.NewHTTPClient(influx.HTTPConfig{
-		Addr:     os.Getenv("INFLUXDB_ADDR"),
+		Addr:     config.InfluxAddr,
-		Username: os.Getenv("INFLUXDB_USERNAME"),
+		Username: config.InfluxUser,
-		Password: os.Getenv("INFLUXDB_PASSWORD"),
+		Password: config.InfluxPass,
 	})
 	if err != nil {
-		panic(err)
+		log.Fatal(err)
 	}
+	log.Println("Polling Unifi Controller, interval:", config.Interval)
+	config.PollUnifiController(infdb)
+}
 
-	log.Printf("Polling Unifi Controller, interval: %+v\n", interval)
+// PollUnifiController runs forever, polling and pushing.
-	for {
+func (c *Config) PollUnifiController(infdb influx.Client) {
-		devices, err := fetch(unifi)
+	ticker := time.NewTicker(c.Interval)
+	for range ticker.C {
+		uniResponse, err := c.GetUnifiClients()
 		if err != nil {
-			log.Println(err)
+			log.Println("GetUnifiClients(unifi):", err)
-		} else {
+			continue
-			bp, _ := influx.NewBatchPoints(influx.BatchPointsConfig{
+		}
-				Database: database,
+		bp, err := influx.NewBatchPoints(influx.BatchPointsConfig{
+			Database: c.InfluxDB,
 		})
-			for _, device := range devices {
+		if err != nil {
-				bp.AddPoint(device.Point())
+			log.Println("influx.NewBatchPoints:", err)
+			continue
 		}
 
-			if err = stats.Write(bp); err != nil {
+		for _, client := range uniResponse.Clients {
-				log.Println(err)
+			bp.AddPoint(client.Point())
 		}
-
+		if err = infdb.Write(bp); err != nil {
-			log.Printf("Logged client state. Devices: %v", len(devices))
+			log.Println("infdb.Write(bp):", err)
+			continue
 		}
-
+		log.Println("Logged client state. Clients:", len(uniResponse.Clients))
-		time.Sleep(interval)
 	}
 }
 
-func fetch(unifi *http.Client) ([]Client, error) {
+// GetConfig parses and returns our configuration data.
-	url := "https://" + os.Getenv("UNIFI_ADDR") + ":" + os.Getenv("UNIFI_PORT") + "/api/s/default/stat/sta"
+func GetConfig() Config {
-	req, err := http.NewRequest("GET", url, nil)
+	// TODO: A real config file.
-	if err != nil {
+	var err error
-		return nil, err
+	config := Config{
+		InfluxAddr: os.Getenv("INFLUXDB_ADDR"),
+		InfluxUser: os.Getenv("INFLUXDB_USERNAME"),
+		InfluxPass: os.Getenv("INFLUXDB_PASSWORD"),
+		InfluxDB:   os.Getenv("INFLUXDB_DATABASE"),
+		UnifiUser:  os.Getenv("UNIFI_USERNAME"),
+		UnifiPass:  os.Getenv("UNIFI_PASSWORD"),
+		UnifiBase:  "https://" + os.Getenv("UNIFI_ADDR") + ":" + os.Getenv("UNIFI_PORT"),
 	}
-	req.Header.Add("Accept", "application/json")
+	if config.Interval, err = time.ParseDuration(os.Getenv("INTERVAL")); err != nil {
-	resp, err := unifi.Do(req)
+		log.Println("Invalid Interval, defaulting to 15 seconds.")
-	if err != nil {
+		config.Interval = time.Duration(time.Second * 15)
-		return nil, err
 	}
-
+	return config
-	defer func() {
-		if err = resp.Body.Close(); err != nil {
-			log.Println("Closing HTTP response:", err)
-		}
-	}()
-	body, _ := ioutil.ReadAll(resp.Body)
-	response := &Response{}
-	err = json.Unmarshal(body, response)
-	if err != nil {
-		return nil, err
-	}
-
-	return response.Data, nil
 }
 
-func login() (*http.Client, error) {
+// GetUnifiClients returns a response full of clients' data from the Unifi Controller.
-	url := "https://" + os.Getenv("UNIFI_ADDR") + ":" + os.Getenv("UNIFI_PORT") + "/api/login"
+func (c *Config) GetUnifiClients() (*ClientResponse, error) {
-	auth := map[string]string{
+	response := &ClientResponse{}
-		"username": os.Getenv("UNIFI_USERNAME"),
+	if req, err := uniRequest("GET", c.UnifiBase+ClientPath, nil); err != nil {
-		"password": os.Getenv("UNIFI_PASSWORD"),
+		return nil, err
+	} else if resp, err := c.uniClient.Do(req); err != nil {
+		return nil, err
+	} else if body, err := ioutil.ReadAll(resp.Body); err != nil {
+		return nil, err
+	} else if err = json.Unmarshal(body, response); err != nil {
+		return nil, err
+	} else if err = resp.Body.Close(); err != nil {
+		log.Println("resp.Body.Close():", err) // Not fatal? Just log it.
 	}
-	transport := &http.Transport{
+	return response, nil
-		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+}
+
+// AuthController creates a http.Client with authenticated cookies.
+// Used to make additional, authenticated requests to the APIs.
+func (c *Config) AuthController() error {
+	json := bytes.NewBufferString(`{"username": "` + c.UnifiUser + `","password": "` + c.UnifiPass + `"}`)
+	jar, err := cookiejar.New(nil)
+	if err != nil {
+		return err
 	}
-	jar, _ := cookiejar.New(nil)
+	c.uniClient = &http.Client{
-	client := &http.Client{
+		Transport: &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}},
-		Transport: transport,
 		Jar:       jar,
 	}
-	json, _ := json.Marshal(auth)
+	if req, err := uniRequest("POST", c.UnifiBase+LoginPath, json); err != nil {
-	params := bytes.NewReader(json)
+		return err
-	req, err := http.NewRequest("POST", url, params)
+	} else if resp, err := c.uniClient.Do(req); err != nil {
-	if err != nil {
+		return err
-		return nil, err
+	} else if resp.StatusCode != http.StatusOK {
+		return errors.New("Error Authenticating with Unifi Controller")
 	}
-
+	return nil
-	req.Header.Add("Accept", "application/json")
+}
-	resp, err := client.Do(req)
+
-	if err != nil {
+func uniRequest(method string, url string, data io.Reader) (*http.Request, error) {
-		return nil, err
+	req, err := http.NewRequest(method, url, data)
-	}
+	if err != nil {
-
+		return nil, err
-	if resp.StatusCode != http.StatusOK {
+	}
-		return nil, errors.New("Error Authenticating with Controller")
+	req.Header.Add("Accept", "application/json")
-	}
+	return req, nil
-
-	return client, nil
 }