126 lines
3.5 KiB
Go
126 lines
3.5 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"io/ioutil"
|
|
"log"
|
|
"os"
|
|
|
|
"github.com/golift/unifi"
|
|
influx "github.com/influxdata/influxdb1-client/v2"
|
|
"github.com/naoina/toml"
|
|
"github.com/pkg/errors"
|
|
flag "github.com/spf13/pflag"
|
|
)
|
|
|
|
func main() {
|
|
u := &UnifiPoller{}
|
|
if u.ParseFlags(os.Args[1:]); u.ShowVer {
|
|
fmt.Printf("unifi-poller v%s\n", Version)
|
|
return // don't run anything else.
|
|
}
|
|
if err := u.GetConfig(); err != nil {
|
|
u.Flag.Usage()
|
|
log.Fatalf("[ERROR] config file '%v': %v", u.ConfigFile, err)
|
|
}
|
|
if err := u.Run(); err != nil {
|
|
log.Fatalln("[ERROR]", err)
|
|
}
|
|
}
|
|
|
|
// ParseFlags runs the parser.
|
|
func (u *UnifiPoller) ParseFlags(args []string) {
|
|
u.Flag = flag.NewFlagSet("unifi-poller", flag.ExitOnError)
|
|
u.Flag.Usage = func() {
|
|
fmt.Println("Usage: unifi-poller [--config=filepath] [--version]")
|
|
u.Flag.PrintDefaults()
|
|
}
|
|
u.Flag.StringVarP(&u.DumpJSON, "dumpjson", "j", "",
|
|
"This debug option prints the json payload for a device and exits.")
|
|
u.Flag.StringVarP(&u.ConfigFile, "config", "c", defaultConfFile, "Poller Config File (TOML Format)")
|
|
u.Flag.BoolVarP(&u.ShowVer, "version", "v", false, "Print the version and exit")
|
|
_ = u.Flag.Parse(args)
|
|
}
|
|
|
|
// GetConfig parses and returns our configuration data.
|
|
func (u *UnifiPoller) GetConfig() error {
|
|
// Preload our defaults.
|
|
u.Config = &Config{
|
|
InfluxURL: defaultInfxURL,
|
|
InfluxUser: defaultInfxUser,
|
|
InfluxPass: defaultInfxPass,
|
|
InfluxDB: defaultInfxDb,
|
|
UnifiUser: defaultUnifUser,
|
|
UnifiPass: os.Getenv("UNIFI_PASSWORD"),
|
|
UnifiBase: defaultUnifURL,
|
|
Interval: Dur{value: defaultInterval},
|
|
Sites: []string{"default"},
|
|
}
|
|
if buf, err := ioutil.ReadFile(u.ConfigFile); err != nil {
|
|
return err
|
|
// This is where the defaults in the config variable are overwritten.
|
|
} else if err := toml.Unmarshal(buf, u.Config); err != nil {
|
|
return err
|
|
}
|
|
if u.DumpJSON != "" {
|
|
u.Quiet = true
|
|
}
|
|
u.Config.Logf("Loaded Configuration: %s", u.ConfigFile)
|
|
return nil
|
|
}
|
|
|
|
// Run invokes all the application logic and routines.
|
|
func (u *UnifiPoller) Run() error {
|
|
c := u.Config
|
|
if u.DumpJSON != "" {
|
|
return c.DumpJSON(u.DumpJSON)
|
|
}
|
|
if log.SetFlags(0); c.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 {
|
|
return err
|
|
}
|
|
infdb, err := c.GetInfluxDB()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
c.PollUnifiController(controller, infdb)
|
|
return nil
|
|
}
|
|
|
|
func (c *Config) GetInfluxDB() (influx.Client, error) {
|
|
infdb, err := influx.NewHTTPClient(influx.HTTPConfig{
|
|
Addr: c.InfluxURL,
|
|
Username: c.InfluxUser,
|
|
Password: c.InfluxPass,
|
|
})
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "influxdb")
|
|
}
|
|
c.Logf("Logging Measurements to InfluxDB at %s as user %s", c.InfluxURL, c.InfluxUser)
|
|
return infdb, nil
|
|
}
|
|
|
|
func (c *Config) GetController() (*unifi.Unifi, error) {
|
|
// Create an authenticated session to the Unifi Controller.
|
|
controller, err := unifi.NewUnifi(c.UnifiUser, c.UnifiPass, c.UnifiBase, c.VerifySSL)
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "unifi controller")
|
|
}
|
|
controller.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.
|
|
}
|
|
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
|
|
}
|
|
c.Logf("Polling Unifi Controller Sites: %v", c.Sites)
|
|
return controller, nil
|
|
}
|