From 6eaf55b06b6d810c1f3cef58d14f022fa9a9c240 Mon Sep 17 00:00:00 2001
From: Chris Grindstaff <chris@gstaff.org>
Date: Fri, 12 Jul 2024 14:24:01 -0400
Subject: [PATCH] feat: Harvest should support StorageGRID credentials script
 with authToken

---
 cmd/collectors/storagegrid/rest/client.go   |  7 +++++++
 pkg/auth/auth.go                            |  9 ++++++---
 pkg/auth/auth_test.go                       | 19 +++++++++++++++++++
 pkg/auth/testdata/get_credentials_authToken |  3 +++
 4 files changed, 35 insertions(+), 3 deletions(-)
 create mode 100755 pkg/auth/testdata/get_credentials_authToken

diff --git a/cmd/collectors/storagegrid/rest/client.go b/cmd/collectors/storagegrid/rest/client.go
index 61ed5341b..0e4122cff 100644
--- a/cmd/collectors/storagegrid/rest/client.go
+++ b/cmd/collectors/storagegrid/rest/client.go
@@ -360,6 +360,13 @@ func (c *Client) fetchTokenWithAuthRetry() error {
 		if err != nil {
 			return err
 		}
+		// If the credential script returns an authToken, use it without re-fetching
+		if pollerAuth.AuthToken != "" {
+			c.token = pollerAuth.AuthToken
+			c.request.Header.Set("Authorization", "Bearer "+c.token)
+			c.Logger.Debug().Msg("Using authToken from credential script")
+			return nil
+		}
 		authB := authBody{
 			Username: pollerAuth.Username,
 			Password: pollerAuth.Password,
diff --git a/pkg/auth/auth.go b/pkg/auth/auth.go
index 52e3aaf3d..618220c46 100644
--- a/pkg/auth/auth.go
+++ b/pkg/auth/auth.go
@@ -123,8 +123,9 @@ func (c *Credentials) fetchCerts(p *conf.Poller) (string, error) {
 }
 
 type ScriptResponse struct {
-	Username string `yaml:"username"`
-	Data     string `yaml:"password"`
+	Username  string `yaml:"username"`
+	Data      string `yaml:"password"`
+	AuthToken string `yaml:"authToken"`
 }
 
 func (c *Credentials) execScript(cmdPath string, kind string, timeout string, e func(ctx context.Context, path string) *exec.Cmd) (ScriptResponse, error) {
@@ -195,7 +196,7 @@ func (c *Credentials) execScript(cmdPath string, kind string, timeout string, e
 			Msg("Failed to parse YAML output. Treating as plain text.")
 	}
 
-	if err == nil && response.Data != "" {
+	if err == nil && (response.Data != "" || response.AuthToken != "") {
 		// If parsing is successful and data is not empty, return the response.
 		// Username is optional, so it's okay if it's not present.
 		return response, nil
@@ -229,6 +230,7 @@ func (c *Credentials) setNextUpdate() {
 type PollerAuth struct {
 	Username             string
 	Password             string
+	AuthToken            string
 	IsCert               bool
 	HasCredentialScript  bool
 	HasCertificateScript bool
@@ -332,6 +334,7 @@ func getPollerAuth(c *Credentials, poller *conf.Poller) (PollerAuth, error) {
 		return PollerAuth{
 			Username:            response.Username,
 			Password:            response.Data,
+			AuthToken:           response.AuthToken,
 			HasCredentialScript: true,
 			Schedule:            poller.CredentialsScript.Schedule,
 			insecureTLS:         insecureTLS,
diff --git a/pkg/auth/auth_test.go b/pkg/auth/auth_test.go
index 6ff1cc595..964ee4b16 100644
--- a/pkg/auth/auth_test.go
+++ b/pkg/auth/auth_test.go
@@ -538,6 +538,22 @@ Pollers:
     username: username
     credentials_script:
       path: testdata/get_credentials_yaml_heredoc
+`,
+		},
+
+		{
+			name:       "credentials_script returns authToken",
+			pollerName: "test",
+			want: PollerAuth{
+				AuthToken:           "abcd",
+				HasCredentialScript: true,
+			},
+			yaml: `
+Pollers:
+  test:
+    addr: a.b.c
+    credentials_script:
+      path: testdata/get_credentials_authToken
 `,
 		},
 	}
@@ -577,6 +593,9 @@ Pollers:
 			if tt.want.Password != got.Password {
 				t.Errorf("got password=[%s], want password=[%s]", got.Password, tt.want.Password)
 			}
+			if tt.want.AuthToken != got.AuthToken {
+				t.Errorf("got authToken=[%s], want authToken=[%s]", got.AuthToken, tt.want.AuthToken)
+			}
 			if tt.want.IsCert != got.IsCert {
 				t.Errorf("got IsCert=[%t], want IsCert=[%t]", got.IsCert, tt.want.IsCert)
 			}
diff --git a/pkg/auth/testdata/get_credentials_authToken b/pkg/auth/testdata/get_credentials_authToken
new file mode 100755
index 000000000..e1487087b
--- /dev/null
+++ b/pkg/auth/testdata/get_credentials_authToken
@@ -0,0 +1,3 @@
+#!/bin/bash
+# Used by pkg/auth/auth_test.go
+echo 'authToken: abcd'
\ No newline at end of file