From cba1111f297e65f0fd14bc98319375fbaec0cf2d Mon Sep 17 00:00:00 2001
From: Cody Lee <cody.lee@datadoghq.com>
Date: Fri, 10 Jan 2025 15:19:30 -0600
Subject: [PATCH] add support for api-key auth

---
 go.mod                      |  2 +-
 go.sum                      |  4 ++--
 pkg/inputunifi/input.go     | 38 +++++++++++++++++++++++++++++--------
 pkg/inputunifi/interface.go |  2 +-
 pkg/inputunifi/updateweb.go |  1 +
 5 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/go.mod b/go.mod
index b7d24a66..6531c615 100644
--- a/go.mod
+++ b/go.mod
@@ -13,7 +13,7 @@ require (
 	github.com/prometheus/common v0.61.0
 	github.com/spf13/pflag v1.0.6-0.20201009195203-85dd5c8bc61c
 	github.com/stretchr/testify v1.10.0
-	github.com/unpoller/unifi/v5 v5.0.7
+	github.com/unpoller/unifi/v5 v5.1.0
 	golang.org/x/crypto v0.31.0
 	golang.org/x/net v0.33.0
 	golang.org/x/term v0.28.0
diff --git a/go.sum b/go.sum
index 4202ac27..1760eaf7 100644
--- a/go.sum
+++ b/go.sum
@@ -75,8 +75,8 @@ github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
 github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-github.com/unpoller/unifi/v5 v5.0.7 h1:Dj5HY2Nhdic4Ygvh2YYW6QKIZjXCSo9IBzVDKaj86Zg=
-github.com/unpoller/unifi/v5 v5.0.7/go.mod h1:G45KRuSH9PFrIUFmDTzWEEM/E/e7GuyXp36AVOfhm7I=
+github.com/unpoller/unifi/v5 v5.1.0 h1:ubX/dugKUXvUlxDupoyPcnevck34dTFDFm3M2n4T75o=
+github.com/unpoller/unifi/v5 v5.1.0/go.mod h1:G45KRuSH9PFrIUFmDTzWEEM/E/e7GuyXp36AVOfhm7I=
 github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
diff --git a/pkg/inputunifi/input.go b/pkg/inputunifi/input.go
index b9a29aa3..870e13f9 100644
--- a/pkg/inputunifi/input.go
+++ b/pkg/inputunifi/input.go
@@ -47,6 +47,7 @@ type Controller struct {
 	CertPaths  []string     `json:"ssl_cert_paths" toml:"ssl_cert_paths" xml:"ssl_cert_path"  yaml:"ssl_cert_paths"`
 	User       string       `json:"user"           toml:"user"           xml:"user"           yaml:"user"`
 	Pass       string       `json:"pass"           toml:"pass"           xml:"pass"           yaml:"pass"`
+	APIKey     string       `json:"api_key"        toml:"api_key"        xml:"api_key"        yaml:"api_key"`
 	URL        string       `json:"url"            toml:"url"            xml:"url"            yaml:"url"`
 	Sites      []string     `json:"sites"          toml:"sites"          xml:"site"           yaml:"sites"`
 	Unifi      *unifi.Unifi `json:"-"              toml:"-"              xml:"-"              yaml:"-"`
@@ -124,6 +125,7 @@ func (u *InputUnifi) getUnifi(c *Controller) error {
 	c.Unifi, err = unifi.NewUnifi(&unifi.Config{
 		User:      c.User,
 		Pass:      c.Pass,
+		APIKey:    c.APIKey,
 		URL:       c.URL,
 		SSLCert:   certs,
 		VerifySSL: *c.VerifySSL,
@@ -255,12 +257,22 @@ func (u *InputUnifi) setDefaults(c *Controller) { //nolint:cyclop
 		c.Pass = u.getPassFromFile(strings.TrimPrefix(c.Pass, "file://"))
 	}
 
-	if c.Pass == "" {
-		c.Pass = defaultPass
+	if strings.HasPrefix(c.APIKey, "file://") {
+		c.APIKey = u.getPassFromFile(strings.TrimPrefix(c.APIKey, "file://"))
 	}
 
-	if c.User == "" {
-		c.User = defaultUser
+	if c.APIKey == "" {
+		if c.Pass == "" {
+			c.Pass = defaultPass
+		}
+
+		if c.User == "" {
+			c.User = defaultUser
+		}
+	} else {
+		// clear out user/pass combo, only use API-key
+		c.User = ""
+		c.Pass = ""
 	}
 
 	if len(c.Sites) == 0 {
@@ -324,12 +336,22 @@ func (u *InputUnifi) setControllerDefaults(c *Controller) *Controller { //nolint
 		c.Pass = u.getPassFromFile(strings.TrimPrefix(c.Pass, "file://"))
 	}
 
-	if c.Pass == "" {
-		c.Pass = u.Default.Pass
+	if strings.HasPrefix(c.APIKey, "file://") {
+		c.APIKey = u.getPassFromFile(strings.TrimPrefix(c.APIKey, "file://"))
 	}
 
-	if c.User == "" {
-		c.User = u.Default.User
+	if c.APIKey == "" {
+		if c.Pass == "" {
+			c.Pass = defaultPass
+		}
+
+		if c.User == "" {
+			c.User = defaultUser
+		}
+	} else {
+		// clear out user/pass combo, only use API-key
+		c.User = ""
+		c.Pass = ""
 	}
 
 	if len(c.Sites) == 0 {
diff --git a/pkg/inputunifi/interface.go b/pkg/inputunifi/interface.go
index 4ffa0420..7b3642e5 100644
--- a/pkg/inputunifi/interface.go
+++ b/pkg/inputunifi/interface.go
@@ -124,7 +124,7 @@ func (u *InputUnifi) logController(c *Controller) {
 		u.Logf("   => Version: %s (%s)", c.Unifi.ServerVersion, c.Unifi.UUID)
 	}
 
-	u.Logf("   => Username: %s (has password: %v)", c.User, c.Pass != "")
+	u.Logf("   => Username: %s (has password: %v) (has api-key: %v)", c.User, c.Pass != "", c.APIKey != "")
 	u.Logf("   => Hash PII %v / Drop PII %v / Poll Sites: %s", *c.HashPII, *c.DropPII, strings.Join(c.Sites, ", "))
 	u.Logf("   => Save Sites %v / Save DPI %v (metrics)", *c.SaveSites, *c.SaveDPI)
 	u.Logf("   => Save Events %v / Save IDs %v (logs)", *c.SaveEvents, *c.SaveIDs)
diff --git a/pkg/inputunifi/updateweb.go b/pkg/inputunifi/updateweb.go
index e9afe7b2..3f07e81d 100644
--- a/pkg/inputunifi/updateweb.go
+++ b/pkg/inputunifi/updateweb.go
@@ -51,6 +51,7 @@ func formatControllers(controllers []*Controller) []*Controller {
 			SaveSites:  c.SaveSites,
 			User:       c.User,
 			Pass:       strconv.FormatBool(c.Pass != ""),
+			APIKey:     strconv.FormatBool(c.APIKey != ""),
 			URL:        c.URL,
 			Sites:      c.Sites,
 			ID:         id,