From d10b521f0ee580a92fb92ebae718820853f9a7e6 Mon Sep 17 00:00:00 2001 From: kxue43 Date: Mon, 15 Dec 2025 23:01:09 -0500 Subject: [PATCH 1/6] doc: Update `README.md`. --- README.md | 1 - 1 file changed, 1 deletion(-) diff --git a/README.md b/README.md index 0fc34aa..b5eb613 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,6 @@ Various CLI programs that aid my own workflows. ``` - `toolkit-show-md` takes a Markdown file, converts it to a GitHub style HTML and displays the HTML in user's default browser. - It's used by [kxue43/showmd-vim-plugin](https://github.com/kxue43/showmd-vim-plugin). ```bash go install github.com/kxue43/cli-toolkit/cmd/toolkit-show-md@latest From 476cee11da97caa7ae89b1255a9df79072e2f852 Mon Sep 17 00:00:00 2001 From: kxue43 Date: Tue, 16 Dec 2025 00:17:53 -0500 Subject: [PATCH 2/6] refactor: Overall structure and style. --- auth/cache.go | 52 ++++++----- auth/command.go | 126 +++++++++++++++----------- auth/command_test.go | 151 -------------------------------- cmd/toolkit-assume-role/main.go | 72 +++++---------- 4 files changed, 126 insertions(+), 275 deletions(-) delete mode 100644 auth/command_test.go diff --git a/auth/cache.go b/auth/cache.go index 5e17446..4b7824d 100644 --- a/auth/cache.go +++ b/auth/cache.go @@ -24,8 +24,8 @@ type ( Version int `json:"Version"` } - cacheSaveRetriever struct { - Logger *log.Logger + Cacher struct { + logger *log.Logger cacheDir string } @@ -37,8 +37,6 @@ type ( cacheFileSlice []*cacheFile ) -const expirationLayout = time.RFC3339 - var ErrInvalidCredential = errors.New("invalid AWS credential") func (cs cacheFileSlice) Len() int { @@ -81,10 +79,12 @@ func DecodeFromFileName(roleArn, fileName string) (ts time.Time, err error) { return ts, nil } -func NewCacheSaveRetriever(logger *log.Logger) (*cacheSaveRetriever, error) { +func NewCacher(logger *log.Logger) *Cacher { home, err := os.UserHomeDir() if err != nil { - return nil, errors.New("could not locate user home directory") + logger.Print("could not locate user home directory") + + return nil } cacheDir := filepath.Join(home, ".aws", "toolkit-cache") @@ -92,29 +92,35 @@ func NewCacheSaveRetriever(logger *log.Logger) (*cacheSaveRetriever, error) { info, err := os.Stat(cacheDir) if os.IsNotExist(err) { if err = os.MkdirAll(cacheDir, 0750); err != nil { - return nil, errors.New("failed to create cache directory") + logger.Print("failed to create cache directory") + + return nil } - return &cacheSaveRetriever{Logger: logger, cacheDir: cacheDir}, nil + return &Cacher{logger: logger, cacheDir: cacheDir} } else if err != nil { - return nil, err + logger.Print(err.Error()) + + return nil } if !info.IsDir() { - return nil, errors.New("cache directory is already a file") + logger.Print("cache directory is already a file") + + return nil } - return &cacheSaveRetriever{Logger: logger, cacheDir: cacheDir}, nil + return &Cacher{logger: logger, cacheDir: cacheDir} } // Save marshals output and saves it to a file whose name is generated according to roleArn and the current timestamp. // On success, it returns the binary contents saved to the file as a byte slice. On failure, if the returned // error wraps ErrInvalidCredential, the AWS credential is invalid. Otherwise only the write operation failed // but the AWS credential is valid. -func (c *cacheSaveRetriever) Save(roleArn string, output *CredentialProcessOutput) (contents []byte, err error) { - ts, err := time.Parse(expirationLayout, output.Expiration) +func (c *Cacher) Save(roleArn string, output *CredentialProcessOutput) (contents []byte, err error) { + ts, err := time.Parse(time.RFC3339, output.Expiration) if err != nil { - return nil, fmt.Errorf("%w: Expiration %q is not of the right format: %w", ErrInvalidCredential, output.Expiration, err) + return nil, fmt.Errorf("%w: expiration %q is not of the right format: %w", ErrInvalidCredential, output.Expiration, err) } fileName := EncodeToFileName(roleArn, ts) @@ -132,14 +138,16 @@ func (c *cacheSaveRetriever) Save(roleArn string, output *CredentialProcessOutpu return contents, nil } -func (c *cacheSaveRetriever) Retrieve(roleArn string) (contents []byte, err error) { +func (c *Cacher) Retrieve(roleArn string) (contents []byte) { max := time.Now().Add(time.Minute * 10) actives := make(cacheFileSlice, 0) pattern := filepath.Join(c.cacheDir, fmt.Sprintf(`%s-*.json`, GetPrefix(roleArn))) cacheFiles, err := filepath.Glob(pattern) if err != nil { - return nil, fmt.Errorf("invalid file globbing pattern: %w", err) + c.logger.Printf("invalid file globbing pattern: %s\n", err) + + return nil } var expiration time.Time @@ -159,7 +167,7 @@ func (c *cacheSaveRetriever) Retrieve(roleArn string) (contents []byte, err erro } if len(actives) == 0 { - return nil, nil + return nil } sort.Sort(actives) @@ -170,14 +178,16 @@ func (c *cacheSaveRetriever) Retrieve(roleArn string) (contents []byte, err erro contents, err = os.ReadFile(actives[0].filePath) if err != nil { - return nil, fmt.Errorf("failed to read active cache file %q: %w", actives[0].filePath, err) + c.logger.Printf("failed to read active cache file %q: %s\n", actives[0].filePath, err) + + return nil } - return contents, nil + return contents } -func (c *cacheSaveRetriever) deleteCacheFile(fullPath string, desc string) { +func (c *Cacher) deleteCacheFile(fullPath string, desc string) { if os.Remove(fullPath) != nil { - c.Logger.Printf("Failed to delete %s cache file %q.\n", desc, fullPath) + c.logger.Printf("Failed to delete %s cache file %q.\n", desc, fullPath) } } diff --git a/auth/command.go b/auth/command.go index 9999f0f..ee76798 100644 --- a/auth/command.go +++ b/auth/command.go @@ -1,8 +1,3 @@ -// Package auth supports performing the AWS CLI credential process. Field tags of the AssumeRoleCmd struct -// are reserved for working with github.com/alecthomas/kong, but the tags are currently unused so that -// the toolkit executable works cross-platform. The toolkit-assume-role executable performs the -// AWS CLI credential process and only works on Linux and macOS because it needs to read from and -// write to /dev/tty. package auth import ( @@ -10,6 +5,7 @@ import ( "context" "encoding/json" "errors" + "flag" "fmt" "io" "log" @@ -23,24 +19,28 @@ import ( type ( AssumeRoleCmd struct { - ci ClientInteractor - cache *cacheSaveRetriever - stsClient *sts.Client - RoleArn string `arg:"" required:"" name:"RoleArn" help:"ARN of the IAM role to assume."` - MFASerial string `required:"" name:"mfa-serial" help:"ARN of the virtual MFA to use when assuming the role."` - Profile string `required:"" name:"profile" help:"Source profile used for assuming the role."` - Region string `name:"region" default:"us-east-1" help:"The regional STS service endpoint to call."` - RoleSessionName string `name:"role-session-name" default:"ToolkitCLI" help:"Role session name."` - DurationSeconds int32 `name:"duration-seconds" default:"3600" help:"Role session duration seconds."` - cacheModeOff bool - } - - ClientInteractor struct { + prompter *Prompter + cacher *Cacher + client *sts.Client + RoleArn string + MFASerial string + Profile string + Region string + RoleSessionName string + DurationSeconds int64 + } + + Prompter struct { + *log.Logger io.ReadWriter } ) -func (c *ClientInteractor) PromptMFAToken() (code string, err error) { +func NewPrompter(tty io.ReadWriter, prefix string, flag int) *Prompter { + return &Prompter{ReadWriter: tty, Logger: log.New(tty, prefix, flag)} +} + +func (c *Prompter) MFAToken() (code string, err error) { _, err = io.WriteString(c, "MFA code: ") if err != nil { return "", fmt.Errorf("failed to prompt for MFA code: %w", err) @@ -56,54 +56,70 @@ func (c *ClientInteractor) PromptMFAToken() (code string, err error) { return string(bytes.TrimSpace(buf[:n])), nil } -func (c *ClientInteractor) NewLogger(prefix string, flag int) *log.Logger { - return log.New(c.ReadWriter, prefix, flag) -} +func (a *AssumeRoleCmd) validate(args []string, logger *log.Logger) { + if a.MFASerial == "" { + logger.Fatal("-mfa-serial is required.") + } -func (a *AssumeRoleCmd) AfterApply(fd io.ReadWriter) error { - cfg, err := config.LoadDefaultConfig(context.TODO(), config.WithSharedConfigProfile(a.Profile), config.WithRegion(a.Region)) - if err != nil { - return fmt.Errorf("failed to load AWS CLI/SDK configuration: %w", err) + if a.Profile == "" { + logger.Fatal("-profile is required.") } - a.stsClient = sts.NewFromConfig(cfg) + if a.DurationSeconds > 14400 { + logger.Fatal("-duration-seconds cannot exceed 14400, i.e. 4 hours.") + } - a.ci = ClientInteractor{fd} - logger := a.ci.NewLogger("toolkit: ", 0) + if len(args) == 0 { + logger.Fatal("The argument is required.") + } - a.cache, err = NewCacheSaveRetriever(logger) - if err != nil { - a.cacheModeOff = true + a.RoleArn = args[0] +} - logger.Printf("cache mode off: %s\n", err) +func (a *AssumeRoleCmd) Init(ctx context.Context, tty io.ReadWriter) { + a.prompter = NewPrompter(tty, "toolkit-assume-role: ", 0) + + a.validate(flag.Args(), a.prompter.Logger) + + cfg, err := config.LoadDefaultConfig(ctx, config.WithSharedConfigProfile(a.Profile), config.WithRegion(a.Region)) + if err != nil { + a.prompter.Fatalf("failed to load AWS CLI/SDK configuration: %s\n", err) } - return nil + a.client = sts.NewFromConfig(cfg) + + a.cacher = NewCacher(a.prompter.Logger) + if a.cacher == nil { + a.prompter.Printf("cache mode off: %s\n", err) + } } -func (a *AssumeRoleCmd) Run(dest io.Writer) (err error) { +func (a *AssumeRoleCmd) Run(ctx context.Context, dest io.Writer) { + var err error + // Output of the AWS CLI credential process. var output []byte - if !a.cacheModeOff { - output, err = a.cache.Retrieve(a.RoleArn) - if err == nil && output != nil { + if a.cacher != nil { + output = a.cacher.Retrieve(a.RoleArn) + if output != nil { _, err = dest.Write(output) - - return err + if err != nil { + a.prompter.Fatalf("failed to write credentials to destination: %s\n", err) + } } } - provider := stscreds.NewAssumeRoleProvider(a.stsClient, a.RoleArn, func(o *stscreds.AssumeRoleOptions) { + provider := stscreds.NewAssumeRoleProvider(a.client, a.RoleArn, func(o *stscreds.AssumeRoleOptions) { o.RoleSessionName = a.RoleSessionName o.Duration = time.Second * time.Duration(a.DurationSeconds) o.SerialNumber = aws.String(a.MFASerial) - o.TokenProvider = a.ci.PromptMFAToken + o.TokenProvider = a.prompter.MFAToken }) - stsCreds, err := provider.Retrieve(context.TODO()) + stsCreds, err := provider.Retrieve(ctx) if err != nil { - return fmt.Errorf("failed to retrieve STS credentials: %w", err) + a.prompter.Fatalf("failed to retrieve STS credentials: %s\n", err) } // structured output @@ -111,27 +127,33 @@ func (a *AssumeRoleCmd) Run(dest io.Writer) (err error) { AccessKeyId: stsCreds.AccessKeyID, SecretAccessKey: stsCreds.SecretAccessKey, SessionToken: stsCreds.SessionToken, - Expiration: stsCreds.Expires.Format(expirationLayout), + Expiration: stsCreds.Expires.Format(time.RFC3339), Version: 1, } - if !a.cacheModeOff { - output, err = a.cache.Save(a.RoleArn, &soutput) + if a.cacher != nil { + output, err = a.cacher.Save(a.RoleArn, &soutput) if errors.Is(err, ErrInvalidCredential) { - return err + a.prompter.Fatal(err.Error()) + } else if err != nil { + a.prompter.Print(err.Error()) } _, err = dest.Write(output) + if err != nil { + a.prompter.Fatalf("failed to write credentials to destination: %s\n", err) + } - return err + return } output, err = json.Marshal(&soutput) if err != nil { - return fmt.Errorf("failed to marshal credential process output: %w", err) + a.prompter.Fatalf("failed to marshal credential process output: %s\n", err) } _, err = dest.Write(output) - - return err + if err != nil { + a.prompter.Fatal(err.Error()) + } } diff --git a/auth/command_test.go b/auth/command_test.go deleted file mode 100644 index 46e1fbb..0000000 --- a/auth/command_test.go +++ /dev/null @@ -1,151 +0,0 @@ -package auth - -import ( - "bytes" - "encoding/json" - "os" - "path/filepath" - "testing" - "time" - - "github.com/aws/aws-sdk-go-v2/service/sts" - "github.com/aws/aws-sdk-go-v2/service/sts/types" - "github.com/awsdocs/aws-doc-sdk-examples/gov2/testtools" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -type ( - MockFileDescriptor struct { - r bytes.Buffer - w bytes.Buffer - } - - HomeDirMocker struct { - TempDir string - home string - } -) - -func (fd *MockFileDescriptor) Read(p []byte) (n int, err error) { - return fd.r.Read(p) -} - -func (fd *MockFileDescriptor) Write(p []byte) (n int, err error) { - return fd.w.Write(p) -} - -func (m *HomeDirMocker) SetUp(t *testing.T) { - t.Helper() - - var err error - - m.TempDir, err = os.MkdirTemp("", "cli-toolkit") - require.NoError(t, err) - - m.home, err = os.UserHomeDir() - require.NoError(t, err) - - err = os.Setenv("HOME", m.TempDir) - require.NoError(t, err) -} - -func (m *HomeDirMocker) TearDown(t *testing.T) { - t.Helper() - - err := os.Setenv("HOME", m.home) - require.NoError(t, err) - - err = os.RemoveAll(m.TempDir) - require.NoError(t, err) -} - -func TestAssumeRoleCmdRun(t *testing.T) { - fd := MockFileDescriptor{} - dest := MockFileDescriptor{} - - hdm := HomeDirMocker{} - - hdm.SetUp(t) - defer hdm.TearDown(t) - - stubber := testtools.NewStubber() - stubbedClient := sts.NewFromConfig(*stubber.SdkConfig) - - ci := ClientInteractor{&fd} - logger := ci.NewLogger("toolkit: ", 0) - - cache, err := NewCacheSaveRetriever(logger) - require.NoError(t, err) - - cmd := AssumeRoleCmd{ - ci: ci, - cache: cache, - stsClient: stubbedClient, - RoleArn: "role-arn", - MFASerial: "mfa-serial", - Profile: "profile", - Region: "us-east-1", - RoleSessionName: "ToolkitCLI", - DurationSeconds: 3600, - } - - token := "123456" - expiration := time.Now() - duration := int32(cmd.DurationSeconds) - - soutput := CredentialProcessOutput{ - AccessKeyId: "access-key-id", - SecretAccessKey: "secret-access-key", - SessionToken: "session-token", - Expiration: expiration.Format(expirationLayout), - Version: 1, - } - - stubber.Add(testtools.Stub{ - OperationName: "AssumeRole", - Input: &sts.AssumeRoleInput{ - DurationSeconds: &duration, - RoleArn: &cmd.RoleArn, - RoleSessionName: &cmd.RoleSessionName, - SerialNumber: &cmd.MFASerial, - TokenCode: &token, - }, - Output: &sts.AssumeRoleOutput{ - Credentials: &types.Credentials{ - AccessKeyId: &soutput.AccessKeyId, - SecretAccessKey: &soutput.SecretAccessKey, - SessionToken: &soutput.SessionToken, - Expiration: &expiration, - }, - }, - Error: nil, - }) - - _, err = fd.r.WriteString(token + "\n") - require.NoError(t, err) - - err = cmd.Run(&dest) - require.NoError(t, err) - - cacheFilePath := filepath.Join(hdm.TempDir, ".aws", "toolkit-cache", EncodeToFileName(cmd.RoleArn, expiration)) - info, err := os.Stat(cacheFilePath) - - require.NoError(t, err) - require.True(t, !info.IsDir()) - - var output CredentialProcessOutput - - contents, err := os.ReadFile(filepath.Clean(cacheFilePath)) - require.NoError(t, err) - - err = json.Unmarshal(contents, &output) - require.NoError(t, err) - - assert.Equal(t, output.AccessKeyId, soutput.AccessKeyId) - assert.Equal(t, output.SecretAccessKey, soutput.SecretAccessKey) - assert.Equal(t, output.SessionToken, soutput.SessionToken) - assert.Equal(t, output.Expiration, soutput.Expiration) - assert.Equal(t, output.Version, soutput.Version) - assert.Equal(t, contents, dest.w.Bytes()) -} diff --git a/cmd/toolkit-assume-role/main.go b/cmd/toolkit-assume-role/main.go index b1439cf..6b9c831 100644 --- a/cmd/toolkit-assume-role/main.go +++ b/cmd/toolkit-assume-role/main.go @@ -1,23 +1,18 @@ package main import ( + "context" "flag" "fmt" - "log" "os" "github.com/kxue43/cli-toolkit/auth" ) var ( - logger = log.New(os.Stderr, "toolkit-assume-role: ", 0) - roleArn string - mfaSerial string - profile string - region string - roleSessionName string - durationSeconds int - helpMsg = `Usage: %s -mfa-serial=STRING -profile=STRING [flags] + cmd = auth.AssumeRoleCmd{} + + helpMsg = `Usage: %s -mfa-serial=STRING -profile=STRING [flags] Run AWS CLI credential process by assuming a role. @@ -29,38 +24,28 @@ Flags: ) func registerFlagsAndHelp() { - flag.StringVar(&mfaSerial, "mfa-serial", "", "ARN of the virtual MFA to use when assuming the role.") - flag.StringVar(&profile, "profile", "", "Source profile used for assuming the role.") - flag.StringVar(®ion, "region", "us-east-1", "The regional STS service endpoint to call.") - flag.StringVar(&roleSessionName, "role-session-name", "ToolkitCLI", "Role session name.") - flag.IntVar(&durationSeconds, "duration-seconds", 3600, "Role session duration seconds.") + flag.StringVar(&cmd.MFASerial, "mfa-serial", "", "ARN of the virtual MFA to use when assuming the role.") + flag.StringVar(&cmd.Profile, "profile", "", "Source profile used for assuming the role.") + flag.StringVar(&cmd.Region, "region", "us-east-1", "The regional STS service endpoint to call.") + flag.StringVar(&cmd.RoleSessionName, "role-session-name", "ToolkitCLI", "Role session name.") + flag.Int64Var(&cmd.DurationSeconds, "duration-seconds", 3600, "Role session duration seconds.") flag.Usage = func() { _, _ = fmt.Fprintf(flag.CommandLine.Output(), helpMsg, os.Args[0]) flag.PrintDefaults() - } -} -func validateInputs() { - if mfaSerial == "" { - logger.Fatalf("-mfa-serial is required.") - } + tty, err := os.OpenFile("/dev/tty", os.O_RDWR|os.O_SYNC, 0600) + if err != nil { + _, _ = fmt.Fprintf(flag.CommandLine.Output(), "\nNot ready to run as credential process: cannot open /dev/tty: %s\n", err) - if profile == "" { - logger.Fatalf("-profile is required.") - } + os.Exit(1) + } - if durationSeconds > 14400 { - logger.Fatal("-duration-seconds cannot exceed 14400, i.e. 4 hours.") - } + defer func() { _ = tty.Close() }() - args := flag.Args() - if len(args) == 0 { - logger.Fatalf("The argument is required.") + _, _ = fmt.Fprint(flag.CommandLine.Output(), "\nReady to run as credential process!\n") } - - roleArn = args[0] } func main() { @@ -68,31 +53,16 @@ func main() { flag.Parse() - validateInputs() - tty, err := os.OpenFile("/dev/tty", os.O_RDWR|os.O_SYNC, 0600) if err != nil { - logger.Fatalf("failed to open /dev/tty: %v", err) + os.Exit(1) } - cmd := auth.AssumeRoleCmd{ - RoleArn: roleArn, - MFASerial: mfaSerial, - Profile: profile, - Region: region, - RoleSessionName: roleSessionName, - DurationSeconds: int32(durationSeconds), - } + defer func() { _ = tty.Close() }() - err = cmd.AfterApply(tty) - if err != nil { - logger.Fatal(err.Error()) - } + ctx := context.Background() - err = cmd.Run(os.Stdout) - if err != nil { - logger.Fatal(err.Error()) - } + cmd.Init(ctx, tty) - defer func() { _ = tty.Close() }() + cmd.Run(ctx, os.Stdout) } From 348718a947db5842abd2f35ad0b9fae85866668b Mon Sep 17 00:00:00 2001 From: kxue43 Date: Tue, 16 Dec 2025 11:29:22 -0500 Subject: [PATCH 3/6] feat: Add encryption to cache file. --- .github/workflows/test-and-lint.yaml | 2 +- .pre-commit-config.yaml | 2 +- auth/cache.go | 37 +++++--- auth/cipher.go | 130 +++++++++++++++++++++++++++ auth/command.go | 9 +- go.mod | 5 +- go.sum | 14 ++- tartufo.toml | 18 ++++ 8 files changed, 197 insertions(+), 20 deletions(-) create mode 100644 auth/cipher.go create mode 100644 tartufo.toml diff --git a/.github/workflows/test-and-lint.yaml b/.github/workflows/test-and-lint.yaml index e6972c0..869e1da 100644 --- a/.github/workflows/test-and-lint.yaml +++ b/.github/workflows/test-and-lint.yaml @@ -42,4 +42,4 @@ jobs: - name: Lint with golangci-lint uses: golangci/golangci-lint-action@v8 with: - version: "v2.6.1" + version: "v2.7.2" diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f6ce9d3..540fe3e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/golangci/golangci-lint - rev: v2.6.1 + rev: v2.7.2 hooks: - id: golangci-lint-fmt - id: golangci-lint-full diff --git a/auth/cache.go b/auth/cache.go index 4b7824d..a449993 100644 --- a/auth/cache.go +++ b/auth/cache.go @@ -27,6 +27,7 @@ type ( Cacher struct { logger *log.Logger cacheDir string + cipher Cipher } cacheFile struct { @@ -58,11 +59,11 @@ func GetPrefix(s string) string { } func EncodeToFileName(roleArn string, ts time.Time) string { - return fmt.Sprintf("%s-%s.json", GetPrefix(roleArn), strconv.FormatInt(ts.Unix(), 10)) + return fmt.Sprintf("%s-%s", GetPrefix(roleArn), strconv.FormatInt(ts.Unix(), 10)) } func DecodeFromFileName(roleArn, fileName string) (ts time.Time, err error) { - regex := regexp.MustCompile(fmt.Sprintf(`^%s-(\d+)\.json$`, GetPrefix(roleArn))) + regex := regexp.MustCompile(fmt.Sprintf(`^%s-(\d+)\$`, GetPrefix(roleArn))) matches := regex.FindStringSubmatch(fileName) if matches == nil { @@ -79,7 +80,7 @@ func DecodeFromFileName(roleArn, fileName string) (ts time.Time, err error) { return ts, nil } -func NewCacher(logger *log.Logger) *Cacher { +func NewCacher(logger *log.Logger, cipher Cipher) *Cacher { home, err := os.UserHomeDir() if err != nil { logger.Print("could not locate user home directory") @@ -97,7 +98,7 @@ func NewCacher(logger *log.Logger) *Cacher { return nil } - return &Cacher{logger: logger, cacheDir: cacheDir} + return &Cacher{logger: logger, cacheDir: cacheDir, cipher: cipher} } else if err != nil { logger.Print(err.Error()) @@ -110,17 +111,13 @@ func NewCacher(logger *log.Logger) *Cacher { return nil } - return &Cacher{logger: logger, cacheDir: cacheDir} + return &Cacher{logger: logger, cacheDir: cacheDir, cipher: cipher} } -// Save marshals output and saves it to a file whose name is generated according to roleArn and the current timestamp. -// On success, it returns the binary contents saved to the file as a byte slice. On failure, if the returned -// error wraps ErrInvalidCredential, the AWS credential is invalid. Otherwise only the write operation failed -// but the AWS credential is valid. func (c *Cacher) Save(roleArn string, output *CredentialProcessOutput) (contents []byte, err error) { ts, err := time.Parse(time.RFC3339, output.Expiration) if err != nil { - return nil, fmt.Errorf("%w: expiration %q is not of the right format: %w", ErrInvalidCredential, output.Expiration, err) + return nil, fmt.Errorf("%w: expiration %q is not of the right format: %s", ErrInvalidCredential, output.Expiration, err.Error()) } fileName := EncodeToFileName(roleArn, ts) @@ -128,11 +125,16 @@ func (c *Cacher) Save(roleArn string, output *CredentialProcessOutput) (contents contents, err = json.Marshal(output) if err != nil { - return nil, fmt.Errorf("%w: failed to serialize CredentialProcessOutput: %w", ErrInvalidCredential, err) + return nil, fmt.Errorf("%w: failed to serialize CredentialProcessOutput: %s", ErrInvalidCredential, err.Error()) } - if err = os.WriteFile(filePath, contents, 0600); err != nil { - return contents, fmt.Errorf("failed to save credentials to cache file: %w", err) + encrypted, err := c.cipher.Encrypt(contents) + if err != nil { + return contents, fmt.Errorf("failed to encrypt cache file: %s", err.Error()) + } + + if err = os.WriteFile(filePath, encrypted, 0600); err != nil { + return contents, fmt.Errorf("failed to write to cache file: %s", err.Error()) } return contents, nil @@ -141,7 +143,7 @@ func (c *Cacher) Save(roleArn string, output *CredentialProcessOutput) (contents func (c *Cacher) Retrieve(roleArn string) (contents []byte) { max := time.Now().Add(time.Minute * 10) actives := make(cacheFileSlice, 0) - pattern := filepath.Join(c.cacheDir, fmt.Sprintf(`%s-*.json`, GetPrefix(roleArn))) + pattern := filepath.Join(c.cacheDir, fmt.Sprintf(`%s-*`, GetPrefix(roleArn))) cacheFiles, err := filepath.Glob(pattern) if err != nil { @@ -183,6 +185,13 @@ func (c *Cacher) Retrieve(roleArn string) (contents []byte) { return nil } + contents, err = c.cipher.Decrypt(contents) + if err != nil { + c.logger.Printf("failed to decrypt cache file %q: %s\n", actives[0].filePath, err) + + return nil + } + return contents } diff --git a/auth/cipher.go b/auth/cipher.go new file mode 100644 index 0000000..1e094fa --- /dev/null +++ b/auth/cipher.go @@ -0,0 +1,130 @@ +package auth + +import ( + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "encoding/base64" + "errors" + "fmt" + "io" + + "github.com/zalando/go-keyring" +) + +type ( + KeyFunc func() ([]byte, error) + + Cipher struct { + key []byte + } +) + +const ( + // 32 bytes for AES-256 + keySize = 32 + + service = "kxue43.toolkit.assume-role" + user = "cache-encryption-key" +) + +var ( + ErrCipher = errors.New("cipher failure") + + fromKeyring KeyFunc = keyringGet +) + +func generateKey(size int) (key []byte, encoded string, err error) { + key = make([]byte, size) + + if _, err = io.ReadFull(rand.Reader, key); err != nil { + return nil, "", fmt.Errorf("%w: failed to generate encryption key: %s", ErrCipher, err.Error()) + } + + return key, base64.StdEncoding.EncodeToString(key), nil +} + +func keyringGet() (key []byte, err error) { + var secret string + + secret, err = keyring.Get(service, user) + if err != nil && !errors.Is(err, keyring.ErrNotFound) { + return nil, fmt.Errorf("%w: failed to retrieve encryption key: secret exists but cannot be read: %s", ErrCipher, err.Error()) + } else if errors.Is(err, keyring.ErrNotFound) { + key, secret, err = generateKey(keySize) + if err != nil { + return nil, err + } + + err = keyring.Set(service, user, secret) + if err != nil { + return nil, fmt.Errorf("%w: failed to save newly generated encryption key: %s", ErrCipher, err.Error()) + } + + return key, nil + } + + key, err = base64.StdEncoding.DecodeString(secret) + if err != nil { + return nil, fmt.Errorf("%w: saved encryption key has been corrupted: %s", ErrCipher, err.Error()) + } + + return key, nil +} + +func NewCipher(fn KeyFunc) (Cipher, error) { + key, err := fn() + if err != nil { + return Cipher{}, fmt.Errorf("%w: failed to get encryption key: %s", ErrCipher, err.Error()) + } + + return Cipher{key: key}, nil +} + +func (c Cipher) Encrypt(plaintext []byte) ([]byte, error) { + block, err := aes.NewCipher(c.key) + if err != nil { + return nil, fmt.Errorf("%w: failed to initialize AES block cipher: %s", ErrCipher, err.Error()) + } + + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, fmt.Errorf("%w: failed to create GCM: %s", ErrCipher, err.Error()) + } + + // The GCM nonce size is fixed at 12 bytes. + nonce := make([]byte, gcm.NonceSize()) + if _, err := io.ReadFull(rand.Reader, nonce); err != nil { + return nil, fmt.Errorf("%w: failed to initialize nonce: %s", ErrCipher, err.Error()) + } + + // The first return value is 'nonce + ciphertext + tag'. + return gcm.Seal(nonce, nonce, plaintext, nil), nil +} + +func (c Cipher) Decrypt(ciphertext []byte) ([]byte, error) { + block, err := aes.NewCipher(c.key) + if err != nil { + return nil, fmt.Errorf("%w: failed to initialize AES block cipher: %s", ErrCipher, err.Error()) + } + + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, fmt.Errorf("%w: failed to create GCM: %s", ErrCipher, err.Error()) + } + + // Extract nonce from the beginning of the ciphertext. + nonceSize := gcm.NonceSize() + if len(ciphertext) < nonceSize { + return nil, fmt.Errorf("%w: ciphertext too short", ErrCipher) + } + + nonce, ciphertextActual := ciphertext[:nonceSize], ciphertext[nonceSize:] + + plaintext, err := gcm.Open(nil, nonce, ciphertextActual, nil) + if err != nil { + return nil, fmt.Errorf("%w: AES-GCM authentication failure, the data have been tampered: %s", ErrCipher, err.Error()) + } + + return plaintext, nil +} diff --git a/auth/command.go b/auth/command.go index ee76798..bd4dbd8 100644 --- a/auth/command.go +++ b/auth/command.go @@ -88,7 +88,14 @@ func (a *AssumeRoleCmd) Init(ctx context.Context, tty io.ReadWriter) { a.client = sts.NewFromConfig(cfg) - a.cacher = NewCacher(a.prompter.Logger) + cipher, err := NewCipher(fromKeyring) + if err != nil { + a.prompter.Printf("failed to create cache cipher, cache mode off: %s\n", err) + + return + } + + a.cacher = NewCacher(a.prompter.Logger, cipher) if a.cacher == nil { a.prompter.Printf("cache mode off: %s\n", err) } diff --git a/go.mod b/go.mod index fc36e3f..3e72638 100644 --- a/go.mod +++ b/go.mod @@ -9,7 +9,6 @@ require ( github.com/aws/aws-sdk-go-v2/config v1.29.9 github.com/aws/aws-sdk-go-v2/credentials v1.17.62 github.com/aws/aws-sdk-go-v2/service/sts v1.33.17 - github.com/awsdocs/aws-doc-sdk-examples/gov2/testtools v0.0.0-20250618174058-74ed3a0539dc github.com/charmbracelet/bubbles v0.21.0 github.com/charmbracelet/bubbletea v1.3.6 github.com/charmbracelet/lipgloss v1.1.0 @@ -18,10 +17,12 @@ require ( github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c github.com/stretchr/testify v1.10.0 github.com/yuin/goldmark v1.7.13 + github.com/zalando/go-keyring v0.2.6 golang.org/x/mod v0.25.0 ) require ( + al.essio.dev/pkg/shellescape v1.5.1 // indirect github.com/atotto/clipboard v0.1.4 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.30 // indirect github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.34 // indirect @@ -37,7 +38,9 @@ require ( github.com/charmbracelet/x/ansi v0.9.3 // indirect github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd // indirect github.com/charmbracelet/x/term v0.2.1 // indirect + github.com/danieljoos/wincred v1.2.2 // indirect github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect + github.com/godbus/dbus/v5 v5.1.0 // indirect github.com/lucasb-eyer/go-colorful v1.2.0 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-localereader v0.0.1 // indirect diff --git a/go.sum b/go.sum index 408dfaf..1640de9 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +al.essio.dev/pkg/shellescape v1.5.1 h1:86HrALUujYS/h+GtqoB26SBEdkWfmMI6FubjXlsXyho= +al.essio.dev/pkg/shellescape v1.5.1/go.mod h1:6sIqp7X2P6mThCQ7twERpZTuigpr6KbZWtls1U8I890= github.com/BurntSushi/toml v1.5.0 h1:W5quZX/G/csjUnuI8SUYlsHs9M38FC7znL0lIO+DvMg= github.com/BurntSushi/toml v1.5.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho= github.com/alecthomas/assert/v2 v2.11.0 h1:2Q9r3ki8+JYXvGsDyBXwH3LcJ+WK5D0gc5E8vS6K3D0= @@ -34,8 +36,6 @@ github.com/aws/aws-sdk-go-v2/service/sts v1.33.17 h1:PZV5W8yk4OtH1JAuhV2PXwwO9v5 github.com/aws/aws-sdk-go-v2/service/sts v1.33.17/go.mod h1:cQnB8CUnxbMU82JvlqjKR2HBOm3fe9pWorWBza6MBJ4= github.com/aws/smithy-go v1.22.2 h1:6D9hW43xKFrRx/tXXfAlIZc4JI+yQe6snnWcQyxSyLQ= github.com/aws/smithy-go v1.22.2/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg= -github.com/awsdocs/aws-doc-sdk-examples/gov2/testtools v0.0.0-20250618174058-74ed3a0539dc h1:5kSZNx8Uj0v/xp9aLCF0K4VEUqKRF42iRGAr5HnW9Y0= -github.com/awsdocs/aws-doc-sdk-examples/gov2/testtools v0.0.0-20250618174058-74ed3a0539dc/go.mod h1:9Oj/8PZn3D5Ftp/Z1QWrIEFE0daERMqfJawL9duHRfc= github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/aymanbagabas/go-udiff v0.2.0 h1:TK0fH4MteXUDspT88n8CKzvK0X9O2xu9yQjWpi6yML8= @@ -56,12 +56,18 @@ github.com/charmbracelet/x/exp/golden v0.0.0-20241011142426-46044092ad91 h1:payR github.com/charmbracelet/x/exp/golden v0.0.0-20241011142426-46044092ad91/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U= github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ= github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg= +github.com/danieljoos/wincred v1.2.2 h1:774zMFJrqaeYCK2W57BgAem/MLi6mtSE47MB6BOJ0i0= +github.com/danieljoos/wincred v1.2.2/go.mod h1:w7w4Utbrz8lqeMbDAK0lkNJUv5sAOkFi7nd/ogr0Uh8= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4= github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM= github.com/goccy/go-yaml v1.18.0 h1:8W7wMFS12Pcas7KU+VVkaiCng+kG8QiFeFwzFb+rwuw= github.com/goccy/go-yaml v1.18.0/go.mod h1:XBurs7gK8ATbW4ZPGKgcbrY1Br56PdM69F7LkFRi1kA= +github.com/godbus/dbus/v5 v5.1.0 h1:4KLkAxT3aOY8Li4FRJe/KvhoNFFxo0m6fNuFUO8QJUk= +github.com/godbus/dbus/v5 v5.1.0/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= +github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4= +github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ= github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY= @@ -85,12 +91,16 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= +github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= 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/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/yuin/goldmark v1.7.13 h1:GPddIs617DnBLFFVJFgpo1aBfe/4xcvMc3SB5t/D0pA= github.com/yuin/goldmark v1.7.13/go.mod h1:ip/1k0VRfGynBgxOz0yCqHrbZXhcjxyuS66Brc7iBKg= +github.com/zalando/go-keyring v0.2.6 h1:r7Yc3+H+Ux0+M72zacZoItR3UDxeWfKTcabvkI8ua9s= +github.com/zalando/go-keyring v0.2.6/go.mod h1:2TCrxYrbUNYfNS/Kgy/LSrkSQzZ5UPVH85RwfczwvcI= golang.org/x/exp v0.0.0-20220909182711-5c715a9e8561 h1:MDc5xs78ZrZr3HMQugiXOAkSZtfTpbJLDr/lwfgO53E= golang.org/x/exp v0.0.0-20220909182711-5c715a9e8561/go.mod h1:cyybsKvd6eL0RnXn6p/Grxp8F5bW7iYuBgsNCOHpMYE= golang.org/x/mod v0.25.0 h1:n7a+ZbQKQA/Ysbyb0/6IbB1H/X41mKgbhfv7AfG/44w= diff --git a/tartufo.toml b/tartufo.toml new file mode 100644 index 0000000..0549652 --- /dev/null +++ b/tartufo.toml @@ -0,0 +1,18 @@ +[tool.tartufo] +repo-path = '.' +regex = true +entropy = true + +exclude-path-patterns = [ + {path-pattern = 'go\.mod$', reason = 'go.mod file'}, + {path-pattern = 'go\.sum$', reason = 'go.sum file'}, + {path-pattern = 'tartufo\.toml$', reason = 'Tartufo config file'}, +] + +exclude-entropy-patterns = [ + # {path-pattern = '', pattern = '', reason = ''}, +] + +exclude-signatures = [ + # {signature = '', reason = ''}, +] From 92c0ad74ed538ad570b6bfb8039d1c04ae6318cb Mon Sep 17 00:00:00 2001 From: kxue43 Date: Tue, 16 Dec 2025 13:46:46 -0500 Subject: [PATCH 4/6] refactor: Return `error` instead of calling `Fatal`. - Pass unit tests. --- .golangci.yaml | 6 +- auth/cache.go | 34 +++--- auth/command.go | 60 +++++------ auth/command_test.go | 186 ++++++++++++++++++++++++++++++++ cmd/toolkit-assume-role/main.go | 25 ++++- go.mod | 1 + go.sum | 2 + 7 files changed, 261 insertions(+), 53 deletions(-) create mode 100644 auth/command_test.go diff --git a/.golangci.yaml b/.golangci.yaml index 0d56815..6b83632 100644 --- a/.golangci.yaml +++ b/.golangci.yaml @@ -112,6 +112,11 @@ linters: - gosec # local static files server doesn't need timeout text: "G114: Use of net/http serve function that has no support for setting timeouts" + - path: '^auth/command_test\.go$' + linters: + - gosec + # unit test only, and we know overflow doesn't happen + text: "G115: integer overflow conversion int64 -> int32" paths: - adhoc/ @@ -128,4 +133,3 @@ formatters: - standard - default - localmodule - diff --git a/auth/cache.go b/auth/cache.go index a449993..0fe887b 100644 --- a/auth/cache.go +++ b/auth/cache.go @@ -38,7 +38,11 @@ type ( cacheFileSlice []*cacheFile ) -var ErrInvalidCredential = errors.New("invalid AWS credential") +var ( + ErrCacheInit = errors.New("cache initialization failure") + ErrCacheSave = errors.New("failed to save cache file") + ErrInvalidCredential = errors.New("invalid AWS credential") +) func (cs cacheFileSlice) Len() int { return len(cs) @@ -80,12 +84,11 @@ func DecodeFromFileName(roleArn, fileName string) (ts time.Time, err error) { return ts, nil } -func NewCacher(logger *log.Logger, cipher Cipher) *Cacher { +// Non-nil returned error wraps [ErrCacheInit]. +func NewCacher(logger *log.Logger, cipher Cipher) (*Cacher, error) { home, err := os.UserHomeDir() if err != nil { - logger.Print("could not locate user home directory") - - return nil + return nil, fmt.Errorf("%w: could not locate user home directory", ErrCacheInit) } cacheDir := filepath.Join(home, ".aws", "toolkit-cache") @@ -93,27 +96,22 @@ func NewCacher(logger *log.Logger, cipher Cipher) *Cacher { info, err := os.Stat(cacheDir) if os.IsNotExist(err) { if err = os.MkdirAll(cacheDir, 0750); err != nil { - logger.Print("failed to create cache directory") - - return nil + return nil, fmt.Errorf("%w: failed to create cache directory", ErrCacheInit) } - return &Cacher{logger: logger, cacheDir: cacheDir, cipher: cipher} + return &Cacher{logger: logger, cacheDir: cacheDir, cipher: cipher}, nil } else if err != nil { - logger.Print(err.Error()) - - return nil + return nil, fmt.Errorf("%w: failed to locate cache directory: %s", ErrCacheInit, err.Error()) } if !info.IsDir() { - logger.Print("cache directory is already a file") - - return nil + return nil, fmt.Errorf("%w: cache directory is already a file", ErrCacheInit) } - return &Cacher{logger: logger, cacheDir: cacheDir, cipher: cipher} + return &Cacher{logger: logger, cacheDir: cacheDir, cipher: cipher}, nil } +// Non-nil returned error wraps [ErrInvalidCredential] or [ErrCacheSave]. func (c *Cacher) Save(roleArn string, output *CredentialProcessOutput) (contents []byte, err error) { ts, err := time.Parse(time.RFC3339, output.Expiration) if err != nil { @@ -130,11 +128,11 @@ func (c *Cacher) Save(roleArn string, output *CredentialProcessOutput) (contents encrypted, err := c.cipher.Encrypt(contents) if err != nil { - return contents, fmt.Errorf("failed to encrypt cache file: %s", err.Error()) + return contents, fmt.Errorf("%w: failed to encrypt before saving: %s", ErrCacheSave, err.Error()) } if err = os.WriteFile(filePath, encrypted, 0600); err != nil { - return contents, fmt.Errorf("failed to write to cache file: %s", err.Error()) + return contents, fmt.Errorf("%w: failed to write to disk: %s", ErrCacheSave, err.Error()) } return contents, nil diff --git a/auth/command.go b/auth/command.go index bd4dbd8..31abf28 100644 --- a/auth/command.go +++ b/auth/command.go @@ -5,14 +5,12 @@ import ( "context" "encoding/json" "errors" - "flag" "fmt" "io" "log" "time" "github.com/aws/aws-sdk-go-v2/aws" - "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/credentials/stscreds" "github.com/aws/aws-sdk-go-v2/service/sts" ) @@ -36,6 +34,10 @@ type ( } ) +var ( + ErrInvalidInput = errors.New("invalid CLI input") +) + func NewPrompter(tty io.ReadWriter, prefix string, flag int) *Prompter { return &Prompter{ReadWriter: tty, Logger: log.New(tty, prefix, flag)} } @@ -56,54 +58,46 @@ func (c *Prompter) MFAToken() (code string, err error) { return string(bytes.TrimSpace(buf[:n])), nil } -func (a *AssumeRoleCmd) validate(args []string, logger *log.Logger) { +// Non-nil returned error wraps [ErrInvalidInput]. +func (a *AssumeRoleCmd) ValidateInputs(args []string) error { if a.MFASerial == "" { - logger.Fatal("-mfa-serial is required.") + return fmt.Errorf("%w: -mfa-serial is required", ErrInvalidInput) } if a.Profile == "" { - logger.Fatal("-profile is required.") + return fmt.Errorf("%w: -profile is required", ErrInvalidInput) } if a.DurationSeconds > 14400 { - logger.Fatal("-duration-seconds cannot exceed 14400, i.e. 4 hours.") + return fmt.Errorf("%w: -duration-seconds cannot exceed 14400, i.e. 4 hours", ErrInvalidInput) } if len(args) == 0 { - logger.Fatal("The argument is required.") + return fmt.Errorf("%w: the argument is required", ErrInvalidInput) } a.RoleArn = args[0] + + return nil } -func (a *AssumeRoleCmd) Init(ctx context.Context, tty io.ReadWriter) { +// Non-nil returned error wraps [ErrCacheInit]. +func (a *AssumeRoleCmd) InitCache(tty io.ReadWriter, cfg aws.Config) error { a.prompter = NewPrompter(tty, "toolkit-assume-role: ", 0) - a.validate(flag.Args(), a.prompter.Logger) - - cfg, err := config.LoadDefaultConfig(ctx, config.WithSharedConfigProfile(a.Profile), config.WithRegion(a.Region)) - if err != nil { - a.prompter.Fatalf("failed to load AWS CLI/SDK configuration: %s\n", err) - } - a.client = sts.NewFromConfig(cfg) cipher, err := NewCipher(fromKeyring) if err != nil { - a.prompter.Printf("failed to create cache cipher, cache mode off: %s\n", err) - - return + return fmt.Errorf("%w: failed to create cache cipher: %s", ErrCacheInit, err.Error()) } - a.cacher = NewCacher(a.prompter.Logger, cipher) - if a.cacher == nil { - a.prompter.Printf("cache mode off: %s\n", err) - } -} + a.cacher, err = NewCacher(a.prompter.Logger, cipher) -func (a *AssumeRoleCmd) Run(ctx context.Context, dest io.Writer) { - var err error + return err +} +func (a *AssumeRoleCmd) Run(ctx context.Context, dest io.Writer) (err error) { // Output of the AWS CLI credential process. var output []byte @@ -112,7 +106,7 @@ func (a *AssumeRoleCmd) Run(ctx context.Context, dest io.Writer) { if output != nil { _, err = dest.Write(output) if err != nil { - a.prompter.Fatalf("failed to write credentials to destination: %s\n", err) + return fmt.Errorf("failed to write credentials to destination: %s", err.Error()) } } } @@ -126,7 +120,7 @@ func (a *AssumeRoleCmd) Run(ctx context.Context, dest io.Writer) { stsCreds, err := provider.Retrieve(ctx) if err != nil { - a.prompter.Fatalf("failed to retrieve STS credentials: %s\n", err) + return fmt.Errorf("failed to retrieve STS credentials: %s", err.Error()) } // structured output @@ -141,26 +135,28 @@ func (a *AssumeRoleCmd) Run(ctx context.Context, dest io.Writer) { if a.cacher != nil { output, err = a.cacher.Save(a.RoleArn, &soutput) if errors.Is(err, ErrInvalidCredential) { - a.prompter.Fatal(err.Error()) + return err } else if err != nil { a.prompter.Print(err.Error()) } _, err = dest.Write(output) if err != nil { - a.prompter.Fatalf("failed to write credentials to destination: %s\n", err) + return fmt.Errorf("failed to write credentials to destination: %s", err.Error()) } - return + return nil } output, err = json.Marshal(&soutput) if err != nil { - a.prompter.Fatalf("failed to marshal credential process output: %s\n", err) + return fmt.Errorf("failed to marshal credential process output: %s", err.Error()) } _, err = dest.Write(output) if err != nil { - a.prompter.Fatal(err.Error()) + return fmt.Errorf("failed to write credentials to destination: %s", err.Error()) } + + return nil } diff --git a/auth/command_test.go b/auth/command_test.go new file mode 100644 index 0000000..9174c17 --- /dev/null +++ b/auth/command_test.go @@ -0,0 +1,186 @@ +package auth + +import ( + "bytes" + "context" + "encoding/json" + "os" + "path/filepath" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/credentials" + "github.com/aws/aws-sdk-go-v2/service/sts" + "github.com/aws/aws-sdk-go-v2/service/sts/types" + "github.com/awsdocs/aws-doc-sdk-examples/gov2/testtools" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type ( + MockFileDescriptor struct { + r bytes.Buffer + w bytes.Buffer + } + + HomeDirMocker struct { + TempDir string + home string + } +) + +func (fd *MockFileDescriptor) Read(p []byte) (n int, err error) { + return fd.r.Read(p) +} + +func (fd *MockFileDescriptor) Write(p []byte) (n int, err error) { + return fd.w.Write(p) +} + +func (m *HomeDirMocker) SetUp(t *testing.T) { + t.Helper() + + var err error + + m.TempDir, err = os.MkdirTemp("", "cli-toolkit") + require.NoError(t, err, "should be able to set up a temp directory for holding cache files during tests") + + m.home, err = os.UserHomeDir() + require.NoError(t, err, "should be able to get user home directory") + + err = os.Setenv("HOME", m.TempDir) + require.NoError(t, err, "should be able to set the HOME environment variable during tests") +} + +func (m *HomeDirMocker) TearDown(t *testing.T) { + t.Helper() + + err := os.Setenv("HOME", m.home) + require.NoError(t, err, "should be able to reset HOME to its original value after tests") + + err = os.RemoveAll(m.TempDir) + require.NoError(t, err, "should be able to remove temp directory after tests") +} + +func TestAssumeRoleCmdRun(t *testing.T) { + fromKeyring = func() ([]byte, error) { + t.Helper() + + key, _, err := generateKey(keySize) + require.NoError(t, err, "should be able to generate a random encryption key") + + return key, nil + } + + defer func() { fromKeyring = keyringGet }() + + tty := MockFileDescriptor{} + dest := MockFileDescriptor{} + + hdm := HomeDirMocker{} + + hdm.SetUp(t) + defer hdm.TearDown(t) + + // Create an AssumeRoleCmd with some fields already filled with valid values. + // We don't test CLI parsing in unit tests. + cmd := AssumeRoleCmd{ + MFASerial: "mfa-serial", + Profile: "profile", + Region: "us-east-1", + RoleSessionName: "ToolkitCLI", + DurationSeconds: 3600, + } + + roleArn := "role-arn" + + token := "123456" + + expiration := time.Now() + + duration := int32(cmd.DurationSeconds) + + soutput := CredentialProcessOutput{ + AccessKeyId: "access-key-id", + SecretAccessKey: "secret-access-key", + SessionToken: "session-token", + Expiration: expiration.Format(time.RFC3339), + Version: 1, + } + + stubber := testtools.NewStubber() + + stubber.Add(testtools.Stub{ + OperationName: "AssumeRole", + Input: &sts.AssumeRoleInput{ + DurationSeconds: &duration, + RoleArn: &roleArn, + RoleSessionName: &cmd.RoleSessionName, + SerialNumber: &cmd.MFASerial, + TokenCode: &token, + }, + Output: &sts.AssumeRoleOutput{ + Credentials: &types.Credentials{ + AccessKeyId: &soutput.AccessKeyId, + SecretAccessKey: &soutput.SecretAccessKey, + SessionToken: &soutput.SessionToken, + Expiration: &expiration, + }, + }, + Error: nil, + }) + + _, err := tty.r.WriteString(token + "\n") + require.NoError(t, err, "should be able to write token to mocked TTY file descriptor") + + ctx := context.Background() + + err = cmd.ValidateInputs([]string{roleArn}) + require.NoError(t, err, "fields of AssumeRoleCmd should validate without error") + + // Use a mocked credentials provider so that we can load config without error during tests. + mockCredsProvider := credentials.NewStaticCredentialsProvider("dummy", "dummy", "dummy") + + cfg, err := config.LoadDefaultConfig(ctx, config.WithCredentialsProvider(mockCredsProvider), config.WithRegion(cmd.Region)) + require.NoError(t, err, "should be able to load config using a mocked credentials provider") + + err = cmd.InitCache(&tty, cfg) + require.NoError(t, err, "should be able to init caching without error") + + // Stub the STS client object. + cmd.client = sts.NewFromConfig(*stubber.SdkConfig) + + err = cmd.Run(ctx, &dest) + require.NoError(t, err, "should be able to run command without error") + + cacheFilePath := filepath.Join(hdm.TempDir, ".aws", "toolkit-cache", EncodeToFileName(roleArn, expiration)) + + info, err := os.Stat(cacheFilePath) + require.NoError(t, err, "should be able to locate the cache file created by the Run method") + + assert.False(t, info.IsDir(), "the cache file created by the Run method should be a regular file, not a directory") + + rawContents, err := os.ReadFile(filepath.Clean(cacheFilePath)) + require.NoError(t, err, "should be able to read cache file raw content without error") + + rawContents, err = cmd.cacher.cipher.Decrypt(rawContents) + require.NoError(t, err, "should be able to decrypt cache file without error") + + var sCachedContents CredentialProcessOutput + + err = json.Unmarshal(rawContents, &sCachedContents) + require.NoError(t, err, "should be able to unmarshal decrypted cache file without error") + + assert.Equal(t, sCachedContents.AccessKeyId, soutput.AccessKeyId, "AccessKeyId from cache file should match STS call result") + + assert.Equal(t, sCachedContents.SecretAccessKey, soutput.SecretAccessKey, "SecretAccessKey from cache file should match STS call result") + + assert.Equal(t, sCachedContents.SessionToken, soutput.SessionToken, "SessionToken from cache file should match STS call result") + + assert.Equal(t, sCachedContents.Expiration, soutput.Expiration, "Expiration from cache file should match STS call result") + + assert.Equal(t, sCachedContents.Version, soutput.Version, "Version from cache file should be the right value of 1") + + assert.Equal(t, rawContents, dest.w.Bytes(), "outputs to stdout should be identical to decrypted cache file contents") +} diff --git a/cmd/toolkit-assume-role/main.go b/cmd/toolkit-assume-role/main.go index 6b9c831..760d927 100644 --- a/cmd/toolkit-assume-role/main.go +++ b/cmd/toolkit-assume-role/main.go @@ -4,8 +4,11 @@ import ( "context" "flag" "fmt" + "log" "os" + "github.com/aws/aws-sdk-go-v2/config" + "github.com/kxue43/cli-toolkit/auth" ) @@ -62,7 +65,25 @@ func main() { ctx := context.Background() - cmd.Init(ctx, tty) + logger := log.New(tty, "toolkit-assume-role: ", 0) + + err = cmd.ValidateInputs(flag.Args()) + if err != nil { + logger.Fatal(err.Error()) + } + + cfg, err := config.LoadDefaultConfig(ctx, config.WithSharedConfigProfile(cmd.Profile), config.WithRegion(cmd.Region)) + if err != nil { + logger.Fatalf("failed to load AWS SDK configuration: %s\n", err) + } + + err = cmd.InitCache(tty, cfg) + if err != nil { + logger.Print(err.Error()) + } - cmd.Run(ctx, os.Stdout) + err = cmd.Run(ctx, os.Stdout) + if err != nil { + logger.Fatal(err.Error()) + } } diff --git a/go.mod b/go.mod index 3e72638..641acc5 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,7 @@ require ( github.com/aws/aws-sdk-go-v2/config v1.29.9 github.com/aws/aws-sdk-go-v2/credentials v1.17.62 github.com/aws/aws-sdk-go-v2/service/sts v1.33.17 + github.com/awsdocs/aws-doc-sdk-examples/gov2/testtools v0.0.0-20251215172815-75f9f7867a88 github.com/charmbracelet/bubbles v0.21.0 github.com/charmbracelet/bubbletea v1.3.6 github.com/charmbracelet/lipgloss v1.1.0 diff --git a/go.sum b/go.sum index 1640de9..800f97a 100644 --- a/go.sum +++ b/go.sum @@ -36,6 +36,8 @@ github.com/aws/aws-sdk-go-v2/service/sts v1.33.17 h1:PZV5W8yk4OtH1JAuhV2PXwwO9v5 github.com/aws/aws-sdk-go-v2/service/sts v1.33.17/go.mod h1:cQnB8CUnxbMU82JvlqjKR2HBOm3fe9pWorWBza6MBJ4= github.com/aws/smithy-go v1.22.2 h1:6D9hW43xKFrRx/tXXfAlIZc4JI+yQe6snnWcQyxSyLQ= github.com/aws/smithy-go v1.22.2/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg= +github.com/awsdocs/aws-doc-sdk-examples/gov2/testtools v0.0.0-20251215172815-75f9f7867a88 h1:EFA5Spdki31BPVKfbS6MwDIwjehQEUNj16w+QGyLeO4= +github.com/awsdocs/aws-doc-sdk-examples/gov2/testtools v0.0.0-20251215172815-75f9f7867a88/go.mod h1:9Oj/8PZn3D5Ftp/Z1QWrIEFE0daERMqfJawL9duHRfc= github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/aymanbagabas/go-udiff v0.2.0 h1:TK0fH4MteXUDspT88n8CKzvK0X9O2xu9yQjWpi6yML8= From 9246a8ef058100af1ccaa75a6501aa141bde620d Mon Sep 17 00:00:00 2001 From: kxue43 Date: Tue, 16 Dec 2025 14:03:47 -0500 Subject: [PATCH 5/6] fix: Return if cache hits. --- auth/cache.go | 5 ++--- auth/command.go | 2 ++ 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/auth/cache.go b/auth/cache.go index 0fe887b..b331a2a 100644 --- a/auth/cache.go +++ b/auth/cache.go @@ -118,14 +118,13 @@ func (c *Cacher) Save(roleArn string, output *CredentialProcessOutput) (contents return nil, fmt.Errorf("%w: expiration %q is not of the right format: %s", ErrInvalidCredential, output.Expiration, err.Error()) } - fileName := EncodeToFileName(roleArn, ts) - filePath := filepath.Join(c.cacheDir, fileName) - contents, err = json.Marshal(output) if err != nil { return nil, fmt.Errorf("%w: failed to serialize CredentialProcessOutput: %s", ErrInvalidCredential, err.Error()) } + filePath := filepath.Join(c.cacheDir, EncodeToFileName(roleArn, ts)) + encrypted, err := c.cipher.Encrypt(contents) if err != nil { return contents, fmt.Errorf("%w: failed to encrypt before saving: %s", ErrCacheSave, err.Error()) diff --git a/auth/command.go b/auth/command.go index 31abf28..477e1c3 100644 --- a/auth/command.go +++ b/auth/command.go @@ -108,6 +108,8 @@ func (a *AssumeRoleCmd) Run(ctx context.Context, dest io.Writer) (err error) { if err != nil { return fmt.Errorf("failed to write credentials to destination: %s", err.Error()) } + + return } } From 655d0dba4c9c66cef2dcf60c0d6d2b8c3f31704b Mon Sep 17 00:00:00 2001 From: kxue43 Date: Tue, 16 Dec 2025 15:10:10 -0500 Subject: [PATCH 6/6] fix: Regexp for cache file name. --- auth/cache.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/auth/cache.go b/auth/cache.go index b331a2a..43a5585 100644 --- a/auth/cache.go +++ b/auth/cache.go @@ -67,7 +67,7 @@ func EncodeToFileName(roleArn string, ts time.Time) string { } func DecodeFromFileName(roleArn, fileName string) (ts time.Time, err error) { - regex := regexp.MustCompile(fmt.Sprintf(`^%s-(\d+)\$`, GetPrefix(roleArn))) + regex := regexp.MustCompile(fmt.Sprintf(`^%s-(\d+)$`, GetPrefix(roleArn))) matches := regex.FindStringSubmatch(fileName) if matches == nil {