From ff2ff0f5d5d1d2a6e4a3c86f7f0a08a636b878e8 Mon Sep 17 00:00:00 2001
From: David Newhall II <david@sleepers.pro>
Date: Fri, 21 Jun 2019 03:58:05 -0700
Subject: [PATCH] Change to newer toml library. Add untested support for xml,
 json and yaml config file formats.

---
 .gitignore                 |  2 ++
 Gopkg.lock                 | 38 ++++++++++++++++++--------------------
 Gopkg.toml                 |  4 ----
 pkg/unifi-poller/config.go | 15 +++++----------
 pkg/unifi-poller/poller.go | 27 +++++++++++++++++++++------
 pkg/unifi-poller/unifi.go  |  4 ++--
 6 files changed, 48 insertions(+), 42 deletions(-)

diff --git a/.gitignore b/.gitignore
index 31237553..ae5e0972 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,9 +1,11 @@
 /up.conf
 /unifi-poller
 /unifi-poller*.gz
+/unifi-poller*.zip
 /unifi-poller*.1
 /unifi-poller*.deb
 /unifi-poller*.rpm
+/unifi-poller.exe
 /unifi-poller.macos
 /unifi-poller.linux
 /unifi-poller.rb
diff --git a/Gopkg.lock b/Gopkg.lock
index 81253253..5147185d 100644
--- a/Gopkg.lock
+++ b/Gopkg.lock
@@ -1,6 +1,14 @@
 # This file is autogenerated, do not edit; changes may be undone by the next 'dep ensure'.
 
 
+[[projects]]
+  digest = "1:9f3b30d9f8e0d7040f729b82dcbc8f0dead820a133b3147ce355fc451f32d761"
+  name = "github.com/BurntSushi/toml"
+  packages = ["."]
+  pruneopts = "UT"
+  revision = "3012a1dbe2e4bd1391d42b32f0577cb7bbc7f005"
+  version = "v0.3.1"
+
 [[projects]]
   digest = "1:28ef1378055e34f154c8efcd8863a3e53a276c58cc7fc0d0a32d6b9eed6f6cfc"
   name = "github.com/golift/unifi"
@@ -21,25 +29,6 @@
   pruneopts = "UT"
   revision = "8ff2fc3824fcb533795f9a2f233275f0bb18d6c5"
 
-[[projects]]
-  digest = "1:b56c589214f01a5601e0821387db484617392d0042f26234bf2da853a2f498a1"
-  name = "github.com/naoina/go-stringutil"
-  packages = ["."]
-  pruneopts = "UT"
-  revision = "6b638e95a32d0c1131db0e7fe83775cbea4a0d0b"
-  version = "v0.1.0"
-
-[[projects]]
-  digest = "1:f58c3d0e46b64878d00652fedba24ee879725191ab919dca7b62586859281c04"
-  name = "github.com/naoina/toml"
-  packages = [
-    ".",
-    "ast",
-  ]
-  pruneopts = "UT"
-  revision = "e6f5723bf2a66af014955e0888881314cf294129"
-  version = "v0.1.1"
-
 [[projects]]
   digest = "1:cf31692c14422fa27c83a05292eb5cbe0fb2775972e8f1f8446a71549bd8980b"
   name = "github.com/pkg/errors"
@@ -56,15 +45,24 @@
   revision = "298182f68c66c05229eb03ac171abe6e309ee79a"
   version = "v1.0.3"
 
+[[projects]]
+  digest = "1:4d2e5a73dc1500038e504a8d78b986630e3626dc027bc030ba5c75da257cdb96"
+  name = "gopkg.in/yaml.v2"
+  packages = ["."]
+  pruneopts = "UT"
+  revision = "51d6538a90f86fe93ac480b35f37b2be17fef232"
+  version = "v2.2.2"
+
 [solve-meta]
   analyzer-name = "dep"
   analyzer-version = 1
   input-imports = [
+    "github.com/BurntSushi/toml",
     "github.com/golift/unifi",
     "github.com/influxdata/influxdb1-client/v2",
-    "github.com/naoina/toml",
     "github.com/pkg/errors",
     "github.com/spf13/pflag",
+    "gopkg.in/yaml.v2",
   ]
   solver-name = "gps-cdcl"
   solver-version = 1
diff --git a/Gopkg.toml b/Gopkg.toml
index 46258e0f..2765dc95 100644
--- a/Gopkg.toml
+++ b/Gopkg.toml
@@ -24,10 +24,6 @@
 #   go-tests = true
 #   unused-packages = true
 
-[[constraint]]
-  name = "github.com/naoina/toml"
-  version = "0.1.1"
-
 [prune]
   go-tests = true
   unused-packages = true
diff --git a/pkg/unifi-poller/config.go b/pkg/unifi-poller/config.go
index 196f7d81..b1804785 100644
--- a/pkg/unifi-poller/config.go
+++ b/pkg/unifi-poller/config.go
@@ -1,7 +1,6 @@
 package unifipoller
 
 import (
-	"strings"
 	"time"
 
 	"github.com/golift/unifi"
@@ -67,14 +66,10 @@ type Config struct {
 }
 
 // Dur is used to UnmarshalTOML into a time.Duration value.
-type Dur struct{ value time.Duration }
+type Dur struct{ time.Duration }
 
-// UnmarshalTOML parses a duration type from a config file.
-func (v *Dur) UnmarshalTOML(data []byte) error {
-	unquoted := strings.Trim(string(data), `"`)
-	dur, err := time.ParseDuration(unquoted)
-	if err == nil {
-		v.value = dur
-	}
-	return err
+// UnmarshalText parses a duration type from a config file.
+func (d *Dur) UnmarshalText(data []byte) (err error) {
+	d.Duration, err = time.ParseDuration(string(data))
+	return
 }
diff --git a/pkg/unifi-poller/poller.go b/pkg/unifi-poller/poller.go
index 80065149..ba75d2ec 100644
--- a/pkg/unifi-poller/poller.go
+++ b/pkg/unifi-poller/poller.go
@@ -1,16 +1,20 @@
 package unifipoller
 
 import (
+	"encoding/json"
+	"encoding/xml"
 	"fmt"
 	"io/ioutil"
 	"log"
 	"os"
+	"strings"
 
+	"github.com/BurntSushi/toml"
 	"github.com/golift/unifi"
 	influx "github.com/influxdata/influxdb1-client/v2"
-	"github.com/naoina/toml"
 	"github.com/pkg/errors"
 	flag "github.com/spf13/pflag"
+	yaml "gopkg.in/yaml.v2"
 )
 
 // ParseFlags runs the parser.
@@ -28,7 +32,7 @@ func (u *UnifiPoller) ParseFlags(args []string) {
 }
 
 // GetConfig parses and returns our configuration data.
-func (u *UnifiPoller) GetConfig() error {
+func (u *UnifiPoller) GetConfig() (err error) {
 	// Preload our defaults.
 	u.Config = &Config{
 		InfluxURL:  defaultInfxURL,
@@ -38,15 +42,26 @@ func (u *UnifiPoller) GetConfig() error {
 		UnifiUser:  defaultUnifUser,
 		UnifiPass:  os.Getenv("UNIFI_PASSWORD"),
 		UnifiBase:  defaultUnifURL,
-		Interval:   Dur{value: defaultInterval},
+		Interval:   Dur{defaultInterval},
 		Sites:      []string{"default"},
 	}
-	if buf, err := ioutil.ReadFile(u.ConfigFile); err != nil {
+	var buf []byte
+	switch buf, err = ioutil.ReadFile(u.ConfigFile); {
+	case 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 {
+	default:
+		err = toml.Unmarshal(buf, u.Config)
+	case strings.HasSuffix(u.ConfigFile, ".json"):
+		err = json.Unmarshal(buf, u.Config)
+	case strings.HasSuffix(u.ConfigFile, ".xml"):
+		err = xml.Unmarshal(buf, u.Config)
+	case strings.HasSuffix(u.ConfigFile, ".yaml"):
+		err = yaml.Unmarshal(buf, u.Config)
+	}
+	if err != nil {
 		return err
 	}
+
 	if u.DumpJSON != "" {
 		u.Quiet = true
 	}
diff --git a/pkg/unifi-poller/unifi.go b/pkg/unifi-poller/unifi.go
index 5bedf73d..19c39c96 100644
--- a/pkg/unifi-poller/unifi.go
+++ b/pkg/unifi-poller/unifi.go
@@ -38,8 +38,8 @@ FIRST:
 
 // PollController runs forever, polling unifi, and pushing to influx.
 func (u *UnifiPoller) PollController() error {
-	log.Println("[INFO] Everything checks out! Poller started, interval:", u.Interval.value)
-	ticker := time.NewTicker(u.Interval.value)
+	log.Println("[INFO] Everything checks out! Poller started, interval:", u.Interval.Round(time.Second))
+	ticker := time.NewTicker(u.Interval.Round(time.Second))
 	var err error
 	for range ticker.C {
 		m := &Metrics{}