Add AWS CLI token cache support

README.md

@@ -5,8 +5,8 @@ This is a command to export the credentials variables to switch a role with MFA.
 Key features:
 
 - Single binary
-- Interoperable config with AWS CLI
-- In-memory credentials (do not write to a file)
+- Interoperable config with AWS CLI (`~/.aws/config`)
+- Interoperable token cache with AWS CLI (`~/.aws/cli/cache`)
 
 
 ## Getting Started

main.go

@@ -12,6 +12,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/aws/endpoints"
 	"github.com/aws/aws-sdk-go-v2/aws/external"
+	"github.com/int128/awsswitch/pkg/tokencache"
 	"golang.org/x/crypto/ssh/terminal"
 )
 
@@ -37,7 +38,6 @@ func run(ctx context.Context, o options) error {
 	if err != nil {
 		return fmt.Errorf("cannot get credentials: %w", err)
 	}
-	log.Printf("you got a valid token until %s", credentials.Expires)
 	fmt.Printf("export AWS_ACCESS_KEY_ID=%s\n", strconv.Quote(credentials.AccessKeyID))
 	fmt.Printf("export AWS_SECRET_ACCESS_KEY=%s\n", strconv.Quote(credentials.SecretAccessKey))
 	fmt.Printf("export AWS_SESSION_TOKEN=%s\n", strconv.Quote(credentials.SessionToken))
@@ -53,13 +53,45 @@ func getCredentials(ctx context.Context, profile string) (*aws.Credentials, erro
 	if err != nil {
 		return nil, fmt.Errorf("cannot load AWS config: %w", err)
 	}
+
+	sc := findSharedConfig(cfg)
+	if sc == nil {
+		credentials, err := cfg.Credentials.Retrieve(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("cannot retrieve credentials: %w", err)
+		}
+		log.Printf("you got a valid token until %s", credentials.Expires)
+		return &credentials, nil
+	}
+
+	cache, err := tokencache.Load(*sc)
+	if err == nil {
+		if !cache.Expired() {
+			log.Printf("you already have a valid token until %s", cache.Expires)
+			return cache, nil
+		}
+		log.Printf("token cache has expired at %s", cache.Expires)
+	}
 	credentials, err := cfg.Credentials.Retrieve(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("cannot retrieve credentials: %w", err)
 	}
+	if err := tokencache.Save(*sc, credentials); err != nil {
+		return nil, fmt.Errorf("cannot save cache: %w", err)
+	}
+	log.Printf("you got a valid token until %s", credentials.Expires)
 	return &credentials, nil
 }
 
+func findSharedConfig(cfg aws.Config) *external.SharedConfig {
+	for _, cs := range cfg.ConfigSources {
+		if sc, ok := cs.(external.SharedConfig); ok {
+			return &sc
+		}
+	}
+	return nil
+}
+
 func readMFACode() (string, error) {
 	if _, err := fmt.Fprint(os.Stderr, mfaCodePrompt); err != nil {
 		return "", fmt.Errorf("cannot write to stderr: %w", err)

pkg/tokencache/content.go

@@ -0,0 +1,67 @@
+package tokencache
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+)
+
+type content struct {
+	Credentials contentCredentials `json:"Credentials"`
+}
+
+type contentCredentials struct {
+	AccessKeyID     string    `json:"AccessKeyId"`
+	SecretAccessKey string    `json:"SecretAccessKey"`
+	SessionToken    string    `json:"SessionToken"`
+	Expiration      time.Time `json:"Expiration"`
+}
+
+func newFromAWSCredentials(credentials aws.Credentials) content {
+	return content{
+		Credentials: contentCredentials{
+			AccessKeyID:     credentials.AccessKeyID,
+			SecretAccessKey: credentials.SecretAccessKey,
+			SessionToken:    credentials.SessionToken,
+			Expiration:      credentials.Expires,
+		},
+	}
+}
+
+func (c *content) awsCredentials() aws.Credentials {
+	return aws.Credentials{
+		AccessKeyID:     c.Credentials.AccessKeyID,
+		SecretAccessKey: c.Credentials.SecretAccessKey,
+		SessionToken:    c.Credentials.SessionToken,
+		Expires:         c.Credentials.Expiration,
+		CanExpire:       true,
+	}
+}
+
+func decodeFrom(name string) (*content, error) {
+	f, err := os.Open(name)
+	if err != nil {
+		return nil, fmt.Errorf("cannot open %s: %w", name, err)
+	}
+	defer f.Close()
+	var c content
+	if err := json.NewDecoder(f).Decode(&c); err != nil {
+		return nil, fmt.Errorf("cannot decode json from %s: %w", name, err)
+	}
+	return &c, nil
+}
+
+func encodeTo(name string, c content) error {
+	f, err := os.OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
+	if err != nil {
+		return fmt.Errorf("cannot create %s: %w", name, err)
+	}
+	defer f.Close()
+	if err := json.NewEncoder(f).Encode(&c); err != nil {
+		return fmt.Errorf("cannot encode json to %s: %w", name, err)
+	}
+	return nil
+}

pkg/tokencache/io.go

@@ -0,0 +1,58 @@
+// Package tokencache provides access to the token cache in ~/.aws/cli/cache.
+// This is interoperable with the token cache of AWS CLI.
+package tokencache
+
+import (
+	"crypto/sha1"
+	"encoding/hex"
+	"fmt"
+	"path/filepath"
+	"strconv"
+	"strings"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/aws/external"
+)
+
+// Load reads the token cache file and returns the credentials.
+func Load(cfg external.SharedConfig) (*aws.Credentials, error) {
+	name := computeFilename(cfg)
+	cache, err := decodeFrom(name)
+	if err != nil {
+		return nil, fmt.Errorf("cannot load the credentials from cache: %w", err)
+	}
+	credentials := cache.awsCredentials()
+	return &credentials, nil
+}
+
+// Save writes the credentials to the token cache file.
+func Save(cfg external.SharedConfig, credentials aws.Credentials) error {
+	name := computeFilename(cfg)
+	if err := encodeTo(name, newFromAWSCredentials(credentials)); err != nil {
+		return fmt.Errorf("cannot save the credentials to %s: %w", name, err)
+	}
+	return nil
+}
+
+func computeFilename(cfg external.SharedConfig) string {
+	key := computeKey(cfg)
+	dir := filepath.Join(filepath.Dir(external.DefaultSharedConfigFilename()), "cli", "cache")
+	return filepath.Join(dir, key+".json")
+}
+
+// computeKey returns a key for the config.
+// This is based on https://github.com/boto/botocore/blob/9f5afe26cc4a2695dcb62db453f9db7c299f3ffc/botocore/credentials.py#L714
+func computeKey(cfg external.SharedConfig) string {
+	var a []string
+	// keys must be sorted
+	if cfg.RoleDurationSeconds != nil {
+		a = append(a, fmt.Sprintf(`"DurationSeconds": %d`, int(cfg.RoleDurationSeconds.Seconds())))
+	}
+	a = append(a, `"RoleArn": `+strconv.Quote(cfg.RoleARN))
+	if cfg.RoleSessionName != "" {
+		a = append(a, `"RoleSessionName": `+strconv.Quote(cfg.RoleSessionName))
+	}
+	a = append(a, `"SerialNumber": `+strconv.Quote(cfg.MFASerial))
+	s := sha1.Sum([]byte(fmt.Sprintf("{%s}", strings.Join(a, ", "))))
+	return hex.EncodeToString(s[:])
+}