This commit is contained in:
David Newhall II 2019-06-13 00:48:16 -07:00
parent 0ebd99b790
commit f4cb6f4699
4 changed files with 69 additions and 66 deletions

View File

@ -3,6 +3,7 @@ package main
import (
"time"
"github.com/golift/unifi"
influx "github.com/influxdata/influxdb1-client/v2"
"github.com/spf13/pflag"
)
@ -27,12 +28,14 @@ type Asset interface {
Points() ([]*influx.Point, error)
}
// UnifiPoller contains the application startup data.
// UnifiPoller contains the application startup data, and auth info for unifi & influx.
type UnifiPoller struct {
ConfigFile string
DumpJSON string
ShowVer bool
Flag *pflag.FlagSet
influx.Client
*unifi.Unifi
*Config
}

View File

@ -9,39 +9,38 @@ import (
"github.com/pkg/errors"
)
// DumpJSON prints raw json from the Unifi Controller.
func (c *Config) DumpJSON(filter string) error {
c.Quiet = true
controller, err := unifi.NewUnifi(c.UnifiUser, c.UnifiPass, c.UnifiBase, c.VerifySSL)
// DumpJSONPayload prints raw json from the Unifi Controller.
func (u *UnifiPoller) DumpJSONPayload() (err error) {
u.Quiet = true
u.Unifi, err = unifi.NewUnifi(u.UnifiUser, u.UnifiPass, u.UnifiBase, u.VerifySSL)
if err != nil {
return err
}
fmt.Fprintln(os.Stderr, "[INFO] Authenticated to Unifi Controller @", c.UnifiBase, "as user", c.UnifiUser)
if err := c.CheckSites(controller); err != nil {
fmt.Fprintln(os.Stderr, "[INFO] Authenticated to Unifi Controller @", u.UnifiBase, "as user", u.UnifiUser)
if err := u.CheckSites(); err != nil {
return err
}
controller.ErrorLog = func(m string, v ...interface{}) {
u.Unifi.ErrorLog = func(m string, v ...interface{}) {
fmt.Fprintf(os.Stderr, "[ERROR] "+m, v...)
} // Log all errors to stderr.
sites, err := filterSites(controller, c.Sites)
switch {
switch sites, err := u.filterSites(u.Sites); {
case err != nil:
return err
case StringInSlice(filter, []string{"d", "device", "devices"}):
return c.DumpDeviceJSON(sites, controller)
case StringInSlice(filter, []string{"client", "clients", "c"}):
return c.DumpClientsJSON(sites, controller)
case StringInSlice(u.DumpJSON, []string{"d", "device", "devices"}):
return u.DumpDeviceJSON(sites)
case StringInSlice(u.DumpJSON, []string{"client", "clients", "c"}):
return u.DumpClientsJSON(sites)
default:
return errors.New("must provide filter: devices, clients")
}
}
// DumpClientsJSON prints the raw json for clients in a Unifi Controller.
func (c *Config) DumpClientsJSON(sites []unifi.Site, controller *unifi.Unifi) error {
func (u *UnifiPoller) DumpClientsJSON(sites []unifi.Site) error {
for _, s := range sites {
path := fmt.Sprintf(unifi.ClientPath, s.Name)
if err := dumpJSON(path, "Client", s, controller); err != nil {
if err := u.dumpJSON(path, "Client", s); err != nil {
return err
}
}
@ -49,22 +48,22 @@ func (c *Config) DumpClientsJSON(sites []unifi.Site, controller *unifi.Unifi) er
}
// DumpDeviceJSON prints the raw json for devices in a Unifi Controller.
func (c *Config) DumpDeviceJSON(sites []unifi.Site, controller *unifi.Unifi) error {
func (u *UnifiPoller) DumpDeviceJSON(sites []unifi.Site) error {
for _, s := range sites {
path := fmt.Sprintf(unifi.DevicePath, s.Name)
if err := dumpJSON(path, "Device", s, controller); err != nil {
if err := u.dumpJSON(path, "Device", s); err != nil {
return err
}
}
return nil
}
func dumpJSON(path, what string, site unifi.Site, controller *unifi.Unifi) error {
req, err := controller.UniReq(path, "")
func (u *UnifiPoller) dumpJSON(path, what string, site unifi.Site) error {
req, err := u.UniReq(path, "")
if err != nil {
return err
}
resp, err := controller.Do(req)
resp, err := u.Do(req)
if err != nil {
return err
}

View File

@ -70,56 +70,56 @@ func (u *UnifiPoller) GetConfig() error {
}
// Run invokes all the application logic and routines.
func (u *UnifiPoller) Run() error {
c := u.Config
func (u *UnifiPoller) Run() (err error) {
if u.DumpJSON != "" {
return c.DumpJSON(u.DumpJSON)
return u.DumpJSONPayload()
}
if log.SetFlags(0); c.Debug {
if log.SetFlags(0); u.Debug {
log.SetFlags(log.Lshortfile | log.Lmicroseconds | log.Ldate)
log.Println("[DEBUG] Debug Logging Enabled")
}
log.Printf("[INFO] Unifi-Poller v%v Starting Up! PID: %d", Version, os.Getpid())
controller, err := c.GetController()
if err != nil {
if err = u.GetUnifi(); err != nil {
return err
}
infdb, err := c.GetInfluxDB()
if err != nil {
if err = u.GetInfluxDB(); err != nil {
return err
}
c.PollUnifiController(controller, infdb)
u.PollController()
return nil
}
func (c *Config) GetInfluxDB() (influx.Client, error) {
infdb, err := influx.NewHTTPClient(influx.HTTPConfig{
Addr: c.InfluxURL,
Username: c.InfluxUser,
Password: c.InfluxPass,
// GetInfluxDB returns an influxdb interface.
func (u *UnifiPoller) GetInfluxDB() (err error) {
u.Client, err = influx.NewHTTPClient(influx.HTTPConfig{
Addr: u.InfluxURL,
Username: u.InfluxUser,
Password: u.InfluxPass,
})
if err != nil {
return nil, errors.Wrap(err, "influxdb")
return errors.Wrap(err, "influxdb")
}
c.Logf("Logging Measurements to InfluxDB at %s as user %s", c.InfluxURL, c.InfluxUser)
return infdb, nil
u.Logf("Logging Measurements to InfluxDB at %s as user %s", u.InfluxURL, u.InfluxUser)
return nil
}
func (c *Config) GetController() (*unifi.Unifi, error) {
// GetUnifi returns a Unifi controller interface.
func (u *UnifiPoller) GetUnifi() (err error) {
// Create an authenticated session to the Unifi Controller.
controller, err := unifi.NewUnifi(c.UnifiUser, c.UnifiPass, c.UnifiBase, c.VerifySSL)
u.Unifi, err = unifi.NewUnifi(u.UnifiUser, u.UnifiPass, u.UnifiBase, u.VerifySSL)
if err != nil {
return nil, errors.Wrap(err, "unifi controller")
return errors.Wrap(err, "unifi controller")
}
controller.ErrorLog = log.Printf // Log all errors.
u.Unifi.ErrorLog = log.Printf // Log all errors.
// Doing it this way allows debug error logs (line numbers, etc)
if c.Debug && !c.Quiet {
controller.DebugLog = log.Printf // Log debug messages.
if u.Debug && !u.Quiet {
u.Unifi.DebugLog = log.Printf // Log debug messages.
}
c.Logf("Authenticated to Unifi Controller at %s as user %s", c.UnifiBase, c.UnifiUser)
if err := c.CheckSites(controller); err != nil {
return nil, err
u.Logf("Authenticated to Unifi Controller at %s as user %s", u.UnifiBase, u.UnifiUser)
if err = u.CheckSites(); err != nil {
return err
}
c.Logf("Polling Unifi Controller Sites: %v", c.Sites)
return controller, nil
u.Logf("Polling Unifi Controller Sites: %v", u.Sites)
return nil
}

View File

@ -11,8 +11,8 @@ import (
)
// CheckSites makes sure the list of provided sites exists on the controller.
func (c *Config) CheckSites(controller *unifi.Unifi) error {
sites, err := controller.GetSites()
func (u *UnifiPoller) CheckSites() error {
sites, err := u.GetSites()
if err != nil {
return err
}
@ -20,12 +20,12 @@ func (c *Config) CheckSites(controller *unifi.Unifi) error {
for _, site := range sites {
msg = append(msg, site.Name+" ("+site.Desc+")")
}
c.Logf("Found %d site(s) on controller: %v", len(msg), strings.Join(msg, ", "))
if StringInSlice("all", c.Sites) {
u.Logf("Found %d site(s) on controller: %v", len(msg), strings.Join(msg, ", "))
if StringInSlice("all", u.Sites) {
return nil
}
FIRST:
for _, s := range c.Sites {
for _, s := range u.Sites {
for _, site := range sites {
if s == site.Name {
continue FIRST
@ -36,28 +36,28 @@ FIRST:
return nil
}
// PollUnifiController runs forever, polling and pushing.
func (c *Config) PollUnifiController(controller *unifi.Unifi, infdb influx.Client) {
log.Println("[INFO] Everything checks out! Poller started, interval:", c.Interval.value)
ticker := time.NewTicker(c.Interval.value)
// PollController runs forever, polling unifi, and pushing to influx.
func (u *UnifiPoller) PollController() {
log.Println("[INFO] Everything checks out! Poller started, interval:", u.Interval.value)
ticker := time.NewTicker(u.Interval.value)
for range ticker.C {
// Get the sites we care about.
sites, err := filterSites(controller, c.Sites)
sites, err := u.filterSites(u.Sites)
if err != nil {
logErrors([]error{err}, "uni.GetSites()")
}
// Get all the points.
clients, err := controller.GetClients(sites)
clients, err := u.GetClients(sites)
if err != nil {
logErrors([]error{err}, "uni.GetClients()")
}
devices, err := controller.GetDevices(sites)
devices, err := u.GetDevices(sites)
if err != nil {
logErrors([]error{err}, "uni.GetDevices()")
}
// Make a new Points Batcher.
bp, err := influx.NewBatchPoints(influx.BatchPointsConfig{Database: c.InfluxDB})
bp, err := influx.NewBatchPoints(influx.BatchPointsConfig{Database: u.InfluxDB})
if err != nil {
logErrors([]error{err}, "influx.NewBatchPoints")
continue
@ -66,9 +66,10 @@ func (c *Config) PollUnifiController(controller *unifi.Unifi, infdb influx.Clien
if errs := batchPoints(devices, clients, bp); errs != nil && hasErr(errs) {
logErrors(errs, "asset.Points()")
}
if err := infdb.Write(bp); err != nil {
if err := u.Write(bp); err != nil {
logErrors([]error{err}, "infdb.Write(bp)")
}
// Talk about the data.
var fieldcount, pointcount int
for _, p := range bp.Points() {
@ -76,7 +77,7 @@ func (c *Config) PollUnifiController(controller *unifi.Unifi, infdb influx.Clien
i, _ := p.Fields()
fieldcount += len(i)
}
c.Logf("Unifi Measurements Recorded. Sites: %d Clients: %d, "+
u.Logf("Unifi Measurements Recorded. Sites: %d Clients: %d, "+
"Wireless APs: %d, Gateways: %d, Switches: %d, Points: %d, Fields: %d",
len(sites), len(clients.UCLs),
len(devices.UAPs), len(devices.USGs), len(devices.USWs), pointcount, fieldcount)
@ -117,8 +118,8 @@ func batchPoints(devices *unifi.Devices, clients *unifi.Clients, bp influx.Batch
// 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()
func (u *UnifiPoller) filterSites(filter []string) ([]unifi.Site, error) {
sites, err := u.GetSites()
if err != nil {
return nil, err
} else if len(filter) < 1 || StringInSlice("all", filter) {