package unifipoller import ( "encoding/json" "encoding/xml" "fmt" "io/ioutil" "log" "os" "strings" "time" "github.com/BurntSushi/toml" influx "github.com/influxdata/influxdb1-client/v2" "github.com/spf13/pflag" "golift.io/unifi" "gopkg.in/yaml.v2" ) // Start begins the application from a CLI. // Parses flags, parses config and executes Run(). func Start() error { log.SetFlags(log.LstdFlags) up := &UnifiPoller{} if up.ParseFlags(os.Args[1:]); up.ShowVer { fmt.Printf("unifi-poller v%s\n", Version) return nil // don't run anything else w/ version request. } if err := up.GetConfig(); err != nil { up.Flag.Usage() return err } return up.Run() } // ParseFlags runs the parser. func (u *UnifiPoller) ParseFlags(args []string) { u.Flag = pflag.NewFlagSet("unifi-poller", pflag.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 a json payload and exits. See man page for more.") 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: Duration{defaultInterval}, Sites: []string{"default"}, Quiet: u.DumpJSON != "", } u.Logf("Loading Configuration File: %s", u.ConfigFile) switch buf, err := ioutil.ReadFile(u.ConfigFile); { case err != nil: return err case strings.Contains(u.ConfigFile, ".json"): return json.Unmarshal(buf, u.Config) case strings.Contains(u.ConfigFile, ".xml"): return xml.Unmarshal(buf, u.Config) case strings.Contains(u.ConfigFile, ".yaml"): return yaml.Unmarshal(buf, u.Config) default: return toml.Unmarshal(buf, u.Config) } } // Run invokes all the application logic and routines. func (u *UnifiPoller) Run() (err error) { if u.DumpJSON != "" { return u.DumpJSONPayload() } if u.Debug { log.SetFlags(log.Lshortfile | log.Lmicroseconds | log.Ldate) u.LogDebugf("Debug Logging Enabled") } log.Printf("[INFO] UniFi Poller v%v Starting Up! PID: %d", Version, os.Getpid()) if err = u.GetUnifi(); err != nil { return err } u.Logf("Polling UniFi Controller at %s v%s as user %s. Sites: %v", u.UnifiBase, u.ServerVersion, u.UnifiUser, u.Sites) if err = u.GetInfluxDB(); err != nil { return err } u.Logf("Logging Measurements to InfluxDB at %s as user %s", u.InfluxURL, u.InfluxUser) switch strings.ToLower(u.Mode) { case "influxlambda", "lambdainflux", "lambda_influx", "influx_lambda": u.LogDebugf("Lambda Mode Enabled") u.LastCheck = time.Now() return u.CollectAndReport() default: return u.PollController() } } // 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 fmt.Errorf("influxdb: %v", err) } return nil } // GetUnifi returns a UniFi controller interface. func (u *UnifiPoller) GetUnifi() (err error) { // Create an authenticated session to the Unifi Controller. u.Unifi, err = unifi.NewUnifi(&unifi.Config{ User: u.UnifiUser, Pass: u.UnifiPass, URL: u.UnifiBase, VerifySSL: u.VerifySSL, ErrorLog: u.LogErrorf, // Log all errors. DebugLog: u.LogDebugf, // Log debug messages. }) if err != nil { return fmt.Errorf("unifi controller: %v", err) } if err := u.CheckSites(); err != nil { return err } return nil }