Cleanup
This commit is contained in:
parent
0ebd99b790
commit
f4cb6f4699
|
|
@ -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
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
Loading…
Reference in New Issue