MWI: flexibility for AWS Roles Anywhere output

lib/tbot/config/service_workload_identity_aws_ra.go

@@ -77,6 +77,19 @@ type WorkloadIdentityAWSRAService struct {
 	// defaults to 1 hour.
 	SessionRenewalInterval time.Duration `yaml:"session_renewal_interval"`
 
+	// CredentialProfileName is the name of the AWS credentials profile to
+	// write to. If unspecified, the profile will be named "default".
+	CredentialProfileName string `yaml:"credential_profile_name,omitempty"`
+
+	// ArtifactName is the name of the artifact to write to. This is the
+	// filename of the file that will be written to the destination. This is
+	// by default "aws_credentials".
+	ArtifactName string `yaml:"artifact_name,omitempty"`
+	// OverwriteCredentialFile is a flag that indicates whether the output
+	// should overwrite the existing credentials file rather than merging with
+	// it.
+	OverwriteCredentialFile bool `yaml:"overwrite_credential_file,omitempty"`
+
 	// EndpointOverride is the endpoint to use for the AWS Roles Anywhere service.
 	// This is designed to be leveraged by tests and unset in production
 	// circumstances.

lib/tbot/config/service_workload_identity_aws_ra_test.go

@@ -30,6 +30,24 @@ func TestWorkloadIdentityAWSRAService_YAML(t *testing.T) {
 	tests := []testYAMLCase[WorkloadIdentityAWSRAService]{
 		{
 			name: "full",
+			in: WorkloadIdentityAWSRAService{
+				Destination: dest,
+				Selector: WorkloadIdentitySelector{
+					Name: "my-workload-identity",
+				},
+				SessionDuration:         time.Minute * 59,
+				SessionRenewalInterval:  time.Minute * 29,
+				RoleARN:                 "arn:aws:iam::123456789012:role/example-role",
+				TrustAnchorARN:          "arn:aws:rolesanywhere:us-east-1:123456789012:trust-anchor/0000000-0000-0000-0000-000000000000",
+				ProfileARN:              "arn:aws:rolesanywhere:us-east-1:123456789012:profile/0000000-0000-0000-0000-00000000000",
+				Region:                  "us-east-1",
+				CredentialProfileName:   "my-profile",
+				ArtifactName:            "my-artifact.toml",
+				OverwriteCredentialFile: true,
+			},
+		},
+		{
+			name: "simple",
 			in: WorkloadIdentityAWSRAService{
 				Destination: dest,
 				Selector: WorkloadIdentitySelector{

lib/tbot/config/testdata/TestWorkloadIdentityAWSRAService_YAML/full.golden

@@ -9,3 +9,6 @@ trust_anchor_arn: arn:aws:rolesanywhere:us-east-1:123456789012:trust-anchor/0000
 region: us-east-1
 session_duration: 59m0s
 session_renewal_interval: 29m0s
+credential_profile_name: my-profile
+artifact_name: my-artifact.toml
+overwrite_credential_file: true

lib/tbot/config/testdata/TestWorkloadIdentityAWSRAService_YAML/simple.golden

@@ -0,0 +1,11 @@
+type: workload-identity-aws-roles-anywhere
+selector:
+  name: my-workload-identity
+destination:
+  type: memory
+role_arn: arn:aws:iam::123456789012:role/example-role
+profile_arn: arn:aws:rolesanywhere:us-east-1:123456789012:profile/0000000-0000-0000-0000-00000000000
+trust_anchor_arn: arn:aws:rolesanywhere:us-east-1:123456789012:trust-anchor/0000000-0000-0000-0000-000000000000
+region: us-east-1
+session_duration: 59m0s
+session_renewal_interval: 29m0s

lib/tbot/service_workload_identity_aws_ra.go

@@ -18,6 +18,7 @@ package tbot
 
 import (
 	"bytes"
+	"cmp"
 	"context"
 	"crypto"
 	"crypto/x509"
@@ -113,7 +114,7 @@ func (s *WorkloadIdentityAWSRAService) generate(ctx context.Context) error {
 		"aws_credentials_expiry", creds.Expiration,
 	)
 
-	return renderAWSCreds(ctx, creds, s.cfg.Destination)
+	return s.renderAWSCreds(ctx, creds)
 }
 
 // exchangeSVID will exchange the X.509 SVID for AWS credentials using the
@@ -224,24 +225,55 @@ func (s *WorkloadIdentityAWSRAService) requestSVID(
 	return x509Credential, privateKey, nil
 }
 
+func loadExistingAWSCredentialFile(
+	ctx context.Context, dest bot.Destination, artifactName string,
+) (*ini.File, error) {
+	// Load the existing credential file if it exists so we can merge with
+	// it.
+	data, err := dest.Read(ctx, artifactName)
+	if err != nil {
+		if trace.IsNotFound(err) {
+			return ini.Empty(), nil
+		}
+		return nil, trace.Wrap(err, "reading existing credentials")
+	}
+
+	f, err := ini.Load(data)
+	if err != nil {
+		return nil, trace.Wrap(err, "parsing existing credentials")
+	}
+	return f, nil
+}
+
 // render will write the AWS credentials to the AWS CLI configuration file.
 // See https://docs.aws.amazon.com/cli/v1/userguide/cli-configure-files.html
-func renderAWSCreds(
-	ctx context.Context, creds *vendoredaws.CredentialProcessOutput, dest bot.Destination,
+func (s *WorkloadIdentityAWSRAService) renderAWSCreds(
+	ctx context.Context,
+	creds *vendoredaws.CredentialProcessOutput,
 ) error {
 	ctx, span := tracer.Start(
 		ctx,
 		"renderAWSCreds",
 	)
 	defer span.End()
 
-	// TODO(noah): At a later date, we can add a mode where we read in and
-	// modify an existing profile within the credentials file - without
-	// overwriting other profiles.
+	artifactName := cmp.Or(s.cfg.ArtifactName, "aws_credentials")
+
 	f := ini.Empty()
-	// "default" is the name of the section that the AWS CLI will use by
+	if !s.cfg.OverwriteCredentialFile {
+		var err error
+		f, err = loadExistingAWSCredentialFile(
+			ctx, s.cfg.Destination, artifactName,
+		)
+		if err != nil {
+			return trace.Wrap(err, "loading existing credentials")
+		}
+	}
+
+	// "default" is the special profile name that the AWS CLI/SDK will read by
 	// default.
-	sec := f.Section("default")
+	profileName := cmp.Or(s.cfg.CredentialProfileName, "default")
+	sec := f.Section(profileName)
 	sec.Key("aws_secret_access_key").SetValue(creds.SecretAccessKey)
 	sec.Key("aws_access_key_id").SetValue(creds.AccessKeyId)
 	sec.Key("aws_session_token").SetValue(creds.SessionToken)
@@ -252,7 +284,7 @@ func renderAWSCreds(
 		return trace.Wrap(err, "writing credentials to buffer")
 	}
 
-	err = dest.Write(ctx, "aws_credentials", b.Bytes())
+	err = s.cfg.Destination.Write(ctx, artifactName, b.Bytes())
 	if err != nil {
 		return trace.Wrap(err, "writing credentials to destination")
 	}

lib/tbot/service_workload_identity_aws_ra_test.go

@@ -49,20 +49,90 @@ func Test_renderAWSCreds(t *testing.T) {
 	}
 	ctx := context.Background()
 
-	dest := &config.DestinationMemory{}
-	require.NoError(t, dest.CheckAndSetDefaults())
-	require.NoError(t, dest.Init(ctx, []string{}))
+	tests := []struct {
+		name         string
+		cfg          *config.WorkloadIdentityAWSRAService
+		artifactName string
+		existingData []byte
+	}{
+		{
+			name:         "normal",
+			cfg:          &config.WorkloadIdentityAWSRAService{},
+			artifactName: "aws_credentials",
+		},
+		{
+			name:         "merge with existing data",
+			cfg:          &config.WorkloadIdentityAWSRAService{},
+			artifactName: "aws_credentials",
+			existingData: []byte(`[foo]
+aws_secret_access_key=existing
+aws_access_key_id=existing
+aws_session_token=existing`),
+		},
+		{
+			name:         "replace with existing data",
+			cfg:          &config.WorkloadIdentityAWSRAService{},
+			artifactName: "aws_credentials",
+			existingData: []byte(`[default]
+aws_secret_access_key=existing
+aws_access_key_id=existing
+aws_session_token=existing`),
+		},
+		{
+			name: "with artifact name override",
+			cfg: &config.WorkloadIdentityAWSRAService{
+				ArtifactName: "foo-xyzzy",
+			},
+			artifactName: "foo-xyzzy",
+		},
+		{
+			name: "with named profile",
+			cfg: &config.WorkloadIdentityAWSRAService{
+				CredentialProfileName: "test-profile",
+			},
+			artifactName: "aws_credentials",
+		},
+		{
+			name: "overwrite existing data",
+			cfg: &config.WorkloadIdentityAWSRAService{
+				CredentialProfileName:   "test-profile",
+				OverwriteCredentialFile: true,
+			},
+			artifactName: "aws_credentials",
+			existingData: []byte(`[foo]
+aws_secret_access_key=existing
+aws_access_key_id=existing
+aws_session_token=existing
+`),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			dest := &config.DestinationMemory{}
+			require.NoError(t, dest.CheckAndSetDefaults())
+			require.NoError(t, dest.Init(ctx, []string{}))
 
-	err := renderAWSCreds(ctx, creds, dest)
-	require.NoError(t, err)
+			if len(tt.existingData) > 0 {
+				require.NoError(t, dest.Write(ctx, tt.artifactName, tt.existingData))
+			}
 
-	got, err := dest.Read(ctx, "aws_credentials")
-	require.NoError(t, err)
+			tt.cfg.Destination = dest
+			svc := &WorkloadIdentityAWSRAService{
+				cfg: tt.cfg,
+			}
 
-	if golden.ShouldSet() {
-		golden.Set(t, got)
+			err := svc.renderAWSCreds(ctx, creds)
+			require.NoError(t, err)
+
+			got, err := dest.Read(ctx, tt.artifactName)
+			require.NoError(t, err)
+
+			if golden.ShouldSet() {
+				golden.Set(t, got)
+			}
+			require.Equal(t, golden.Get(t), got)
+		})
 	}
-	require.Equal(t, golden.Get(t), got)
 }
 
 type mockCreateSessionInputBody struct {

lib/tbot/testdata/Test_renderAWSCreds/merge_with_existing_data.golden

@@ -0,0 +1,9 @@
+[foo]
+aws_secret_access_key=existing
+aws_access_key_id=existing
+aws_session_token=existing
+
+[default]
+aws_secret_access_key=wJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLESAK
+aws_access_key_id=AKIAIOSFODNN7EXAMPLEAKID
+aws_session_token=AQoDYXdzEJrtyWJ4NjK7PiEXAMPLEST

lib/tbot/testdata/Test_renderAWSCreds/overwrite_existing_data.golden

@@ -0,0 +1,4 @@
+[test-profile]
+aws_secret_access_key=wJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLESAK
+aws_access_key_id=AKIAIOSFODNN7EXAMPLEAKID
+aws_session_token=AQoDYXdzEJrtyWJ4NjK7PiEXAMPLEST

lib/tbot/testdata/Test_renderAWSCreds/replace_with_existing_data.golden

@@ -0,0 +1,4 @@
+[default]
+aws_secret_access_key=wJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLESAK
+aws_access_key_id=AKIAIOSFODNN7EXAMPLEAKID
+aws_session_token=AQoDYXdzEJrtyWJ4NjK7PiEXAMPLEST

lib/tbot/testdata/Test_renderAWSCreds/with_artifact_name_override.golden

@@ -0,0 +1,4 @@
+[default]
+aws_secret_access_key=wJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLESAK
+aws_access_key_id=AKIAIOSFODNN7EXAMPLEAKID
+aws_session_token=AQoDYXdzEJrtyWJ4NjK7PiEXAMPLEST

lib/tbot/testdata/Test_renderAWSCreds/with_named_profile.golden

@@ -0,0 +1,4 @@
+[test-profile]
+aws_secret_access_key=wJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLESAK
+aws_access_key_id=AKIAIOSFODNN7EXAMPLEAKID
+aws_session_token=AQoDYXdzEJrtyWJ4NjK7PiEXAMPLEST