Merge pull request #471 from lacework/afiune/ALLY-546/cli/cache

feat: Introducing a caching mechanism

api/auth.go

@@ -23,6 +23,8 @@ import (
 	"time"
 
 	"go.uber.org/zap"
+
+	"github.com/lacework/go-sdk/internal/format"
 )
 
 const DefaultTokenExpiryTime = 3600
@@ -46,7 +48,7 @@ func WithApiKeys(id, secret string) Option {
 
 		c.log.Debug("setting up auth",
 			zap.String("key", id),
-			zap.String("secret", secret),
+			zap.String("secret", format.Secret(4, secret)),
 		)
 		c.auth.keyID = id
 		c.auth.secret = secret
@@ -70,9 +72,23 @@ func WithTokenFromKeys(id, secret string) Option {
 // WithToken sets the token used to authenticate the API requests
 func WithToken(token string) Option {
 	return clientFunc(func(c *Client) error {
-		c.log.Debug("setting up auth", zap.String("token", token))
+		c.log.Debug("setting up auth", zap.String("token", format.Secret(4, token)))
+		c.auth.token = token
+		c.auth.expiresAt = time.Now().UTC().Add(DefaultTokenExpiryTime * time.Second)
+		return nil
+	})
+}
+
+// WithTokenAndExpiration sets the token used to authenticate the API requests
+// and additionally configures the expiration of the token
+func WithTokenAndExpiration(token string, expiration time.Time) Option {
+	return clientFunc(func(c *Client) error {
+		c.log.Debug("setting up auth",
+			zap.String("token", format.Secret(4, token)),
+			zap.Time("expires_at", expiration),
+		)
 		c.auth.token = token
-		c.auth.expiresAt = time.Now().Add(DefaultTokenExpiryTime * time.Second)
+		c.auth.expiresAt = expiration.UTC()
 		return nil
 	})
 }
@@ -81,15 +97,14 @@ func WithToken(token string) Option {
 func WithExpirationTime(t int) Option {
 	return clientFunc(func(c *Client) error {
 		c.log.Debug("setting up auth", zap.Int("expiration", t))
-
 		c.auth.expiration = t
-		c.auth.expiresAt = time.Now().Add(time.Duration(t) * time.Second)
+		c.auth.expiresAt = time.Now().UTC().Add(time.Duration(t) * time.Second)
 		return nil
 	})
 }
 
 func (c *Client) TokenExpired() bool {
-	return time.Until(c.auth.expiresAt) <= 0
+	return c.auth.expiresAt.Sub(time.Now().UTC()) <= 0
 }
 
 // GenerateToken generates a new access token
@@ -118,20 +133,23 @@ func (c *Client) GenerateToken() (*TokenData, error) {
 		defer res.Body.Close()
 	default:
 		// we default to v1
-		var tokenResponse TokenResponse
-		res, err := c.DoDecoder(request, &tokenResponse)
+		var tokenV1 TokenV1Response
+		res, err := c.DoDecoder(request, &tokenV1)
 		if err != nil {
 			return nil, err
 		}
 		defer res.Body.Close()
 
-		tokenData.Token = tokenResponse.Token()
-		tokenData.ExpiresAt = tokenResponse.ExpiresAt()
+		tokenData.Token = tokenV1.Token()
+		tokenData.ExpiresAt = tokenV1.ExpiresAt()
 	}
 
-	c.log.Debug("storing token", zap.Reflect("data", tokenData))
+	c.log.Debug("storing token",
+		zap.String("token", format.Secret(4, tokenData.Token)),
+		zap.Time("expires_at", tokenData.ExpiresAt),
+	)
 	c.auth.token = tokenData.Token
-	c.auth.expiresAt, err = time.Parse(time.RFC3339, tokenData.ExpiresAt)
+	c.auth.expiresAt = tokenData.ExpiresAt
 	if err != nil {
 		c.log.Error("failed to parse token expiration response", zap.Error(err))
 	}
@@ -142,46 +160,53 @@ func (c *Client) GenerateToken() (*TokenData, error) {
 func (c *Client) GenerateTokenWithKeys(keyID, secretKey string) (*TokenData, error) {
 	c.log.Debug("setting up auth",
 		zap.String("key", keyID),
-		zap.String("secret", secretKey),
+		zap.String("secret", format.Secret(4, secretKey)),
 	)
 	c.auth.keyID = keyID
 	c.auth.secret = secretKey
 	return c.GenerateToken()
 }
 
-type TokenResponse struct {
-	Data    []TokenData `json:"data"`
-	Ok      bool        `json:"ok"`
-	Message string      `json:"message"`
+type tokenRequest struct {
+	KeyID      string `json:"keyId"`
+	ExpiryTime int    `json:"expiryTime"`
 }
 
-func (tr TokenResponse) Token() string {
-	if len(tr.Data) > 0 {
-		// @afiune how do we handle cases where there is more than one token
-		return tr.Data[0].Token
-	}
+// APIv2
+type TokenData struct {
+	ExpiresAt time.Time `json:"expiresAt"`
+	Token     string    `json:"token"`
+}
 
-	return ""
+// APIv1
+type TokenV1Data struct {
+	ExpiresAt string `json:"expiresAt"`
+	Token     string `json:"token"`
 }
 
-func (tr TokenResponse) ExpiresAt() string {
-	if len(tr.Data) > 0 {
-		// @afiune how do we handle cases where there is more than one token
-		expiresAtTime, err := time.Parse("Jan 02 2006 15:04", tr.Data[0].ExpiresAt)
-		if err == nil {
-			return expiresAtTime.Format(time.RFC3339)
-		}
+type TokenV1Response struct {
+	Data    []TokenV1Data `json:"data"`
+	Ok      bool          `json:"ok"`
+	Message string        `json:"message"`
+}
+
+// Soon-To-Be-Deprecated
+func (v1 TokenV1Response) Token() string {
+	if len(v1.Data) > 0 {
+		return v1.Data[0].Token
 	}
 
 	return ""
 }
 
-type TokenData struct {
-	ExpiresAt string `json:"expiresAt"`
-	Token     string `json:"token"`
-}
+// Soon-To-Be-Deprecated
+func (v1 TokenV1Response) ExpiresAt() time.Time {
+	if len(v1.Data) > 0 {
+		expiresAtTime, err := time.Parse("Jan 02 2006 15:04", v1.Data[0].ExpiresAt)
+		if err == nil {
+			return expiresAtTime
+		}
+	}
 
-type tokenRequest struct {
-	KeyID      string `json:"keyId"`
-	ExpiryTime int    `json:"expiryTime"`
+	return time.Now().UTC()
 }

api/auth_test.go

@@ -129,3 +129,21 @@ func TestGenerateTokenErrorKeysMissing(t *testing.T) {
 		}
 	}
 }
+
+func TestGenerateTokenV1(t *testing.T) {
+	fakeServer := lacework.MockServer()
+	fakeServer.MockToken("EXPECTED_TOKEN")
+	defer fakeServer.Close()
+
+	c, err := api.NewClient("t",
+		api.WithURL(fakeServer.URL()),
+		api.WithApiKeys("KEY", "SECRET"),
+	)
+
+	assert.Nil(t, err)
+
+	tokenData, err := c.GenerateToken()
+	if assert.Nil(t, err) {
+		assert.Equal(t, "EXPECTED_TOKEN", tokenData.Token)
+	}
+}

api/callbacks.go

@@ -0,0 +1,39 @@
+//
+// Author:: Salim Afiune Maya (<[email protected]>)
+// Copyright:: Copyright 2021, Lacework Inc.
+// License:: Apache License, Version 2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+package api
+
+import "net/http"
+
+type LifecycleCallbacks struct {
+	// RequestCallback is a function that will be executed after every client request
+	RequestCallback func(int, http.Header) error
+
+	// TokenExpiredCallback  is a function that the consumer can configure
+	// into the client so that it is run when the token expired
+	TokenExpiredCallback func() error
+}
+
+// WithLifecycleCallbacks will configure the lifecycle callback functions
+func WithLifecycleCallbacks(callbacks LifecycleCallbacks) Option {
+	return clientFunc(func(c *Client) error {
+		c.log.Debug("setting up client callbacks")
+		c.callbacks = callbacks
+		return nil
+	})
+}

api/client.go

@@ -43,6 +43,7 @@ type Client struct {
 	c          *http.Client
 	log        *zap.Logger
 	headers    map[string]string
+	callbacks  LifecycleCallbacks
 
 	Account         *AccountService
 	Agents          *AgentsService

api/http.go

@@ -52,8 +52,15 @@ func (c *Client) NewRequest(method string, apiURL string, body io.Reader) (*http
 	if apiURL == apiTokens {
 		headers["X-LW-UAKS"] = c.auth.secret
 	} else {
-		// verify that the client has a token or token is not expired, if not, try to generate one
+		// verify that the client has a token or token is not expired,
+		// if not, try to generate one
 		if c.auth.token == "" || c.TokenExpired() {
+			// run token expired callback
+			if c.callbacks.TokenExpiredCallback != nil && c.TokenExpired() {
+				if err := c.callbacks.TokenExpiredCallback(); err != nil {
+					c.log.Info("token expired callback failure", zap.String("error", err.Error()))
+				}
+			}
 			if _, err = c.GenerateToken(); err != nil {
 				return nil, err
 			}
@@ -166,6 +173,14 @@ func (c *Client) Do(req *http.Request) (*http.Response, error) {
 			zap.String("body", c.httpResponseBodySniffer(response)),
 		)
 	}
+
+	// run request callback
+	if call := c.callbacks.RequestCallback; call != nil {
+		if err := call(response.StatusCode, response.Header); err != nil {
+			c.log.Info("request callback failure", zap.String("error", err.Error()))
+		}
+	}
+
 	return response, err
 }
 

cli/cmd/access_token.go

@@ -67,6 +67,7 @@ func generateAccessToken(_ *cobra.Command, args []string) error {
 		// if the duration is different from the default,
 		// regenerate the lacework api client
 		client, err := api.NewClient(cli.Account,
+			api.WithApiV2(),
 			api.WithLogLevel(cli.LogLevel),
 			api.WithExpirationTime(durationSeconds),
 			api.WithHeader("User-Agent", fmt.Sprintf("Command-Line/%s", Version)),
@@ -81,6 +82,15 @@ func generateAccessToken(_ *cobra.Command, args []string) error {
 		}
 	}
 
+	// cache new token
+	err = cli.Cache.Write("token", structToString(response))
+	if err != nil {
+		cli.Log.Warnw("unable to write token in cache",
+			"feature", "cache",
+			"error", err.Error(),
+		)
+	}
+
 	if cli.JSONOutput() {
 		return cli.OutputJSON(response)
 	}