From ecfce4b9dec574bb4add283f181ab50969114615 Mon Sep 17 00:00:00 2001
From: Rafael Grigorian <rafael@grigorian.org>
Date: Mon, 1 Jul 2024 23:25:20 -0500
Subject: [PATCH] Refreshing client token instead of reissuing another one

---
 sdk/credentials/session.go | 25 +++++++++++++++++
 sdk/tui/tui.go             | 57 ++++++++++++++++++++------------------
 2 files changed, 55 insertions(+), 27 deletions(-)

diff --git a/sdk/credentials/session.go b/sdk/credentials/session.go
index d7fca2b..09c5f87 100644
--- a/sdk/credentials/session.go
+++ b/sdk/credentials/session.go
@@ -326,6 +326,31 @@ func (s *Session) WaitForToken(deviceCode string) error {
 	return nil
 }
 
+func (s *Session) RefreshToken() error {
+	options := ssooidc.Options{Region: s.Region}
+	client := ssooidc.New(options)
+	token, err := client.CreateToken(context.TODO(), &ssooidc.CreateTokenInput{
+		ClientId:     aws.String(s.ClientCredentials.ClientId),
+		ClientSecret: aws.String(s.ClientCredentials.ClientSecret),
+		RefreshToken: aws.String(s.ClientToken.RefreshToken),
+		GrantType:    aws.String("refresh_token"),
+	})
+	if err != nil {
+		return err
+	}
+	s.ClientToken = &ClientToken{
+		AccessToken:           aws.ToString(token.AccessToken),
+		ClientId:              s.ClientCredentials.ClientId,
+		ClientSecret:          s.ClientCredentials.ClientSecret,
+		ExpiresAt:             time.Now().Add(time.Duration(token.ExpiresIn) * time.Second).UTC(),
+		RefreshToken:          aws.ToString(token.RefreshToken),
+		Region:                s.Region,
+		RegistrationExpiresAt: s.ClientCredentials.ExpiresAt,
+		StartUrl:              s.StartUrl,
+	}
+	return nil
+}
+
 func (s *Session) GetAccounts() (Accounts, error) {
 	accounts := Accounts{}
 	options := sso.Options{Region: s.Region}
diff --git a/sdk/tui/tui.go b/sdk/tui/tui.go
index 782daf3..e5cae54 100644
--- a/sdk/tui/tui.go
+++ b/sdk/tui/tui.go
@@ -23,37 +23,40 @@ var (
 )
 
 func ClientLogin(session *credentials.Session) error {
-	if err := session.RegisterClient(); err != nil {
-		return err
-	}
-	userCode, deviceCode, url, urlFull, err := session.StartDeviceAuthorization()
-	if err != nil {
-		return err
-	}
-	yellow := color.ToForeground(YellowColor).Decorator()
-	gray := color.ToForeground(LightGrayColor).Decorator()
-	title := TitleStyle.Decorator()
-	DefaultStyle.Printfln("")
-	DefaultStyle.Printfln("%s %s", title("SSO Session:      "), gray(session.Name))
-	DefaultStyle.Printfln("%s %s", title("SSO Start URL:    "), gray(session.StartUrl))
-	DefaultStyle.Printfln("%s %s", title("Authorization URL:"), gray(url))
-	DefaultStyle.Printfln("%s %s", title("Device Code:      "), yellow(userCode))
-	DefaultStyle.Printfln("")
-	DefaultStyle.Printf("Waiting for authorization to complete...")
-	err = browser.OpenURL(urlFull)
-	if err != nil {
+	if session.ClientCredentials.IsExpired() {
+		if err := session.RegisterClient(); err != nil {
+			return err
+		}
+		userCode, deviceCode, url, urlFull, err := session.StartDeviceAuthorization()
+		if err != nil {
+			return err
+		}
+		yellow := color.ToForeground(YellowColor).Decorator()
+		gray := color.ToForeground(LightGrayColor).Decorator()
+		title := TitleStyle.Decorator()
+		DefaultStyle.Printfln("")
+		DefaultStyle.Printfln("%s %s", title("SSO Session:      "), gray(session.Name))
+		DefaultStyle.Printfln("%s %s", title("SSO Start URL:    "), gray(session.StartUrl))
+		DefaultStyle.Printfln("%s %s", title("Authorization URL:"), gray(url))
+		DefaultStyle.Printfln("%s %s", title("Device Code:      "), yellow(userCode))
+		DefaultStyle.Printfln("")
+		DefaultStyle.Printf("Waiting for authorization to complete...")
+		err = browser.OpenURL(urlFull)
+		if err != nil {
+			ansi.MoveCursorUp(6)
+			ansi.ClearDown()
+			return err
+		}
+		err = session.WaitForToken(deviceCode)
 		ansi.MoveCursorUp(6)
 		ansi.ClearDown()
+		if err != nil {
+			return err
+		}
+	} else if err := session.RefreshToken(); err != nil {
 		return err
 	}
-	err = session.WaitForToken(deviceCode)
-	ansi.MoveCursorUp(6)
-	ansi.ClearDown()
-	if err != nil {
-		return err
-	}
-	err = session.Save()
-	if err != nil {
+	if err := session.Save(); err != nil {
 		return err
 	}
 	return nil