Restructure the code.

This commit is contained in:
DN2 2018-04-20 23:42:03 -07:00
parent d0f1899311
commit 863ed5c75a
2 changed files with 231 additions and 158 deletions

View File

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

388
main.go
View File

@ -5,25 +5,48 @@ import (
"crypto/tls" "crypto/tls"
"encoding/json" "encoding/json"
"errors" "errors"
"io"
"io/ioutil" "io/ioutil"
"log" "log"
"net/http" "net/http"
"net/http/cookiejar" "net/http/cookiejar"
"os" "os"
"strconv"
"time" "time"
influx "github.com/influxdata/influxdb/client/v2" 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"
Meta struct { // ClientPath is Unifi Clients API Path
Rc string ClientPath = "/api/s/default/stat/sta"
} )
// ClientResponse marshalls the payload from the controller.
type ClientResponse struct {
Clients []Client `json:"data"`
Meta struct {
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. // DpiStat is for deep packet inspection stats.
// Does not seem to exist in Unifi 5.7.20.
type DpiStat struct { type DpiStat struct {
App int64 App int64
Cat int64 Cat int64
@ -35,94 +58,136 @@ type DpiStat struct {
// Client defines all the data a connected-network client contains. // Client defines all the data a connected-network client contains.
type Client struct { type Client struct {
ID string `json:"_id"` ID string `json:"_id"`
IsGuestByUAP bool `json:"_is_guest_by_uap"` IsGuestByUAP bool `json:"_is_guest_by_uap"`
IsGuestByUGW bool `json:"_is_guest_by_ugw"` IsGuestByUGW bool `json:"_is_guest_by_ugw"`
LastSeenByUAP int64 `json:"_last_seen_by_uap"` IsGuestByUSW bool `json:"_is_guest_by_usw"`
LastSeenByUGW int64 `json:"_last_seen_by_ugw"` LastSeenByUAP int64 `json:"_last_seen_by_uap"`
UptimeByUAP int64 `json:"_uptime_by_uap"` LastSeenByUGW int64 `json:"_last_seen_by_ugw"`
UptimeByUGW int64 `json:"_uptime_by_ugw"` LastSeenByUSW int64 `json:"_last_seen_by_usw"`
ApMac string `json:"ap_mac"` UptimeByUAP int64 `json:"_uptime_by_uap"`
AssocTime int64 `json:"assoc_time"` UptimeByUGW int64 `json:"_uptime_by_ugw"`
Authorized bool UptimeByUSW int64 `json:"_uptime_by_usw"`
Bssid string ApMac string `json:"ap_mac"`
BytesR int64 `json:"bytes-r"` AssocTime int64 `json:"assoc_time"`
Ccq int64 Authorized bool `json:"authorized"`
Channel int64 Bssid string `json:"bssid"`
BytesR int64 `json:"bytes-r"`
Ccq int64 `json:"ccq"`
Channel int `json:"channel"`
DpiStats []DpiStat `json:"dpi_stats"` DpiStats []DpiStat `json:"dpi_stats"`
DpiStatsLastUpdated int64 `json:"dpi_stats_last_updated"` DpiStatsLastUpdated int64 `json:"dpi_stats_last_updated"`
Essid string Essid string `json:"essid"`
FirstSeen int64 `json:"first_seen"` FirstSeen int64 `json:"first_seen"`
FixedIP string `json:"fixed_ip"` FixedIP string `json:"fixed_ip"`
Hostname string `json:"hostname"` Hostname string `json:"hostname"`
GwMac string `json:"gw_mac"` GwMac string `json:"gw_mac"`
IdleTime int64 `json:"idle_time"` IdleTime int64 `json:"idle_time"`
IP string IP string `json:"ip"`
IsGuest bool `json:"is_guest"` Is11R bool `json:"is_11r"`
IsWired bool `json:"is_wired"` IsGuest bool `json:"is_guest"`
LastSeen int64 `json:"last_seen"` IsWired bool `json:"is_wired"`
LatestAssocTime int64 `json:"latest_assoc_time"` LastSeen int64 `json:"last_seen"`
Mac string LatestAssocTime int64 `json:"latest_assoc_time"`
Name string Mac string `json:"mac"`
Network string Name string `json:"name"`
NetworkID string `json:"network_id"` Network string `json:"network"`
Noise int64 NetworkID string `json:"network_id"`
Oui string Noise int64 `json:"noise"`
PowersaveEnabled bool `json:"powersave_enabled"` Note string `json:"note"`
QosPolicyApplied bool `json:"qos_policy_applied"` Noted bool `json:"noted"`
Radio string Oui string `json:"oui"`
RadioName string `json:"radio_name"` PowersaveEnabled bool `json:"powersave_enabled"`
RadioProto string `json:"radio_proto"` QosPolicyApplied bool `json:"qos_policy_applied"`
RoamCount int64 `json:"roam_count"` Radio string `json:"radio"`
Rssi int64 RadioName string `json:"radio_name"`
RxBytes int64 `json:"rx_bytes"` RadioProto string `json:"radio_proto"`
RxBytesR int64 `json:"rx_bytes-r"` RoamCount int64 `json:"roam_count"`
RxPackets int64 `json:"rx_packets"` Rssi int64 `json:"rssi"`
RxRate int64 `json:"rx_rate"` RxBytes int64 `json:"rx_bytes"`
Signal int64 RxBytesR int64 `json:"rx_bytes-r"`
SiteID string `json:"site_id"` RxPackets int64 `json:"rx_packets"`
TxBytes int64 `json:"tx_bytes"` RxRate int64 `json:"rx_rate"`
TxBytesR int64 `json:"tx_bytes-r"` Signal int64 `json:"signal"`
TxPackets int64 `json:"tx_packets"` SiteID string `json:"site_id"`
TxPower int64 `json:"tx_power"` SwDepth int `json:"sw_depth"`
TxRate int64 `json:"tx_rate"` SwMac string `json:"sw_mac"`
Uptime int64 SwPort int `json:"sw_port"`
UserID string `json:"user_id"` TxBytes int64 `json:"tx_bytes"`
Vlan int64 TxBytesR int64 `json:"tx_bytes-r"`
TxPackets int64 `json:"tx_packets"`
TxPower int64 `json:"tx_power"`
TxRate int64 `json:"tx_rate"`
Uptime int64 `json:"uptime"`
UserID string `json:"user_id"`
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 { 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{ tags := map[string]string{
"mac": c.Mac, "id": c.ID,
"user_id": c.UserID, "mac": c.Mac,
"site_id": c.SiteID, "user_id": c.UserID,
"ip": c.IP, "site_id": c.SiteID,
"essid": c.Essid, "ip": c.IP,
"network": c.Network, "fixed_ip": c.FixedIP,
"ap_mac": c.ApMac, "essid": c.Essid,
"name": c.Name, "bssid": c.Bssid,
"hostname": c.Hostname, "network": c.Network,
"radio_name": c.RadioName, "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{}{ 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_uap": c.LastSeenByUAP,
"last_seen_by_ugw": c.LastSeenByUGW, "last_seen_by_ugw": c.LastSeenByUGW,
"last_seen_by_usw": c.LastSeenByUSW,
"uptime_by_uap": c.UptimeByUAP, "uptime_by_uap": c.UptimeByUAP,
"uptime_by_ugw": c.UptimeByUGW, "uptime_by_ugw": c.UptimeByUGW,
"uptime_by_usw": c.UptimeByUSW,
"assoc_time": c.AssocTime, "assoc_time": c.AssocTime,
"bytes_r": c.BytesR, "bytes_r": c.BytesR,
"ccq": c.Ccq, "ccq": c.Ccq,
"channel": c.Channel,
"dpi_stats_last_updated": c.DpiStatsLastUpdated,
"first_seen": c.FirstSeen, "first_seen": c.FirstSeen,
"idle_time": c.IdleTime, "idle_time": c.IdleTime,
"last_seen": c.LastSeen, "last_seen": c.LastSeen,
"latest_assoc_time": c.LatestAssocTime, "latest_assoc_time": c.LatestAssocTime,
"noise": c.Noise, "noise": c.Noise,
"note": c.Note,
"roam_count": c.RoamCount, "roam_count": c.RoamCount,
"rssi": c.Rssi, "rssi": c.Rssi,
"rx_bytes": c.RxBytes, "rx_bytes": c.RxBytes,
@ -136,7 +201,6 @@ func (c Client) Point() *influx.Point {
"tx_power": c.TxPower, "tx_power": c.TxPower,
"tx_rate": c.TxRate, "tx_rate": c.TxRate,
"uptime": c.Uptime, "uptime": c.Uptime,
"vlan": c.Vlan,
} }
pt, err := influx.NewPoint("client_state", tags, fields, time.Now()) pt, err := influx.NewPoint("client_state", tags, fields, time.Now())
@ -149,108 +213,116 @@ func (c Client) Point() *influx.Point {
} }
func main() { 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!") 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 { 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 { 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 { })
bp.AddPoint(device.Point()) if err != nil {
} log.Println("influx.NewBatchPoints:", err)
continue
if err = stats.Write(bp); err != nil {
log.Println(err)
}
log.Printf("Logged client state. Devices: %v", len(devices))
} }
time.Sleep(interval) for _, client := range uniResponse.Clients {
bp.AddPoint(client.Point())
}
if err = infdb.Write(bp); err != nil {
log.Println("infdb.Write(bp):", err)
continue
}
log.Println("Logged client state. Clients:", len(uniResponse.Clients))
} }
} }
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, 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
} }