[v18] Fix tsh app login when setting up access for AWS using Roles Anywhere integration

lib/aws/awsconfigfile/awsconfigfile.go

@@ -17,7 +17,6 @@
 package awsconfigfile
 
 import (
-	"cmp"
 	"os"
 	"path/filepath"
 	"strings"
@@ -26,6 +25,10 @@ import (
 	ini "gopkg.in/ini.v1"
 )
 
+const (
+	ownershipComment = "Do not edit. Section managed by Teleport."
+)
+
 // AWSConfigFilePath returns the path to the AWS configuration file.
 func AWSConfigFilePath() (string, error) {
 	// See https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html
@@ -43,19 +46,19 @@ func AWSConfigFilePath() (string, error) {
 
 // SetDefaultProfileCredentialProcess sets the credential_process for the default profile.
 // File is created if it does not exist.
-func SetDefaultProfileCredentialProcess(configFilePath, sectionComment, credentialProcess string) error {
+func SetDefaultProfileCredentialProcess(configFilePath, credentialProcess string) error {
 	const sectionName = "default"
-	return trace.Wrap(addCredentialProcessToSection(configFilePath, sectionName, sectionComment, credentialProcess))
+	return trace.Wrap(addCredentialProcessToSection(configFilePath, sectionName, credentialProcess))
 }
 
 // UpsertProfileCredentialProcess sets the credential_process for the profile with name profileName.
 // File is created if it does not exist.
-func UpsertProfileCredentialProcess(configFilePath, profileName, sectionComment, credentialProcess string) error {
+func UpsertProfileCredentialProcess(configFilePath, profileName, credentialProcess string) error {
 	sectionName := "profile " + profileName
-	return trace.Wrap(addCredentialProcessToSection(configFilePath, sectionName, sectionComment, credentialProcess))
+	return trace.Wrap(addCredentialProcessToSection(configFilePath, sectionName, credentialProcess))
 }
 
-func addCredentialProcessToSection(configFilePath, sectionName, sectionComment, credentialProcessCommand string) error {
+func addCredentialProcessToSection(configFilePath, sectionName, credentialProcessCommand string) error {
 	iniFile, err := ini.LoadSources(ini.LoadOptions{
 		AllowNestedValues: true, // Allow AWS-like nested values. Docs: http://docs.aws.amazon.com/cli/latest/topic/config-vars.html#nested-values
 		Loose:             true, // Allow non-existing files. ini.SaveTo will create the file if it does not exist.
@@ -64,16 +67,24 @@ func addCredentialProcessToSection(configFilePath, sectionName, sectionComment,
 		return trace.Wrap(err)
 	}
 
-	if !iniFile.HasSection(sectionName) {
-		iniFile.NewSection(sectionName)
-	}
+	var section *ini.Section
 
-	section := iniFile.Section(sectionName)
-	if cmp.Or(section.Comment, sectionComment) != sectionComment {
-		return trace.BadParameter("%s: section %q is not managed by teleport, remove the section and try again", configFilePath, sectionName)
+	switch {
+	case iniFile.HasSection(sectionName):
+		section = iniFile.Section(sectionName)
+
+		if !strings.Contains(section.Comment, ownershipComment) {
+			return trace.BadParameter("%s: section %q is not managed by Teleport, remove the section and try again", configFilePath, sectionName)
+		}
+
+	default:
+		section, err = iniFile.NewSection(sectionName)
+		if err != nil {
+			return trace.Wrap(err)
+		}
 	}
 
-	section.Comment = sectionComment
+	section.Comment = ownershipComment
 	_, err = section.NewKey("credential_process", credentialProcessCommand)
 	if err != nil {
 		return trace.Wrap(err)
@@ -91,33 +102,46 @@ func addCredentialProcessToSection(configFilePath, sectionName, sectionComment,
 	return trace.Wrap(iniFile.SaveTo(configFilePath))
 }
 
-// RemoveCredentialProcessByComment removes the credential_process key on sections that have a specific section comment.
-func RemoveCredentialProcessByComment(configFilePath, sectionComment string) error {
-	if !strings.HasPrefix(sectionComment, "; ") {
-		sectionComment = "; " + sectionComment
-	}
+// RemoveTeleportManagedProfile removes the credential_process key on sections that have a specific section comment.
+func RemoveTeleportManagedProfile(configFilePath, profile string) error {
+	sectionName := "profile " + profile
 
-	compareExactlyFn := func(comment string) bool {
-		return sectionComment == comment
+	iniFile, err := ini.LoadSources(ini.LoadOptions{
+		AllowNestedValues: true,  // Allow AWS-like nested values. Docs: http://docs.aws.amazon.com/cli/latest/topic/config-vars.html#nested-values
+		Loose:             false, // If file does not exist, then there's nothing to be removed.
+	}, configFilePath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			// Ignore non-existing file.
+			return nil
+		}
+
+		return trace.Wrap(err)
 	}
 
-	return removeCredentialProcess(configFilePath, compareExactlyFn)
-}
+	configFileChanged := false
+	if iniFile.HasSection(sectionName) {
+		section := iniFile.Section(sectionName)
 
-// RemoveCredentialProcessByCommentPrefix removes the credential_process key on sections that have a specific section comment prefix.
-func RemoveCredentialProcessByCommentPrefix(configFilePath, sectionComment string) error {
-	if !strings.HasPrefix(sectionComment, "; ") {
-		sectionComment = "; " + sectionComment
-	}
+		if !strings.Contains(section.Comment, ownershipComment) {
+			return trace.BadParameter("%s: section %q is not managed by Teleport, remove the section manually and try again", configFilePath, sectionName)
+		}
 
-	comparePrefixFn := func(comment string) bool {
-		return strings.HasPrefix(comment, sectionComment)
+		if strings.Contains(section.Comment, ownershipComment) {
+			iniFile.DeleteSection(section.Name())
+			configFileChanged = true
+		}
+	}
+	// No need to save the file if no sections were changed.
+	if !configFileChanged {
+		return nil
 	}
-	return removeCredentialProcess(configFilePath, comparePrefixFn)
+
+	return trace.Wrap(iniFile.SaveTo(configFilePath))
 }
 
-// RemoveCredentialProcessByComment removes the credential_process key on sections that have a specific section comment.
-func removeCredentialProcess(configFilePath string, matchCommentFn func(string) bool) error {
+// RemoveAllTeleportManagedProfiles removes all the profiles managed by Teleport.
+func RemoveAllTeleportManagedProfiles(configFilePath string) error {
 	iniFile, err := ini.LoadSources(ini.LoadOptions{
 		AllowNestedValues: true,  // Allow AWS-like nested values. Docs: http://docs.aws.amazon.com/cli/latest/topic/config-vars.html#nested-values
 		Loose:             false, // If file does not exist, then there's nothing to be removed.
@@ -133,21 +157,10 @@ func removeCredentialProcess(configFilePath string, matchCommentFn func(string)
 
 	sectionChanged := false
 	for _, section := range iniFile.Sections() {
-		if !matchCommentFn(section.Comment) {
-			continue
-		}
-
-		if !section.HasKey("credential_process") {
-			continue
-		}
-
-		section.DeleteKey("credential_process")
-		if len(section.Keys()) > 0 {
-			return trace.BadParameter("%s: section %q contains other keys, remove the section and try again", configFilePath, section.Name())
+		if strings.Contains(section.Comment, ownershipComment) {
+			iniFile.DeleteSection(section.Name())
+			sectionChanged = true
 		}
-		iniFile.DeleteSection(section.Name())
-
-		sectionChanged = true
 	}
 
 	// No need to save the file if no sections were changed.