Use a struct to capture inherited code

Author: Chris Gilmer

cmd/add_profile.go

@@ -12,7 +12,6 @@ import (
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
-	"gopkg.in/go-playground/validator.v9"
 	"gopkg.in/ini.v1"
 )
 
@@ -57,54 +56,20 @@ func AddProfileCheckConfig(v *viper.Viper) error {
 	return nil
 }
 
-// AddProfileConfig holds information for the AWS user being configured by this script
+// AddProfileConfig holds information for the AWS profile configured by this script
 type AddProfileConfig struct {
-	Logger *log.Logger
-	Config *vault.ConfigFile
-
-	IAMRole   string
-	Partition string
-	Region    string
-	Output    string
-
-	BaseProfileName    string
-	AWSProfileName     string
-	AWSProfileAccounts []string
-	AWSProfiles        []vault.ProfileSection
-	MFASerial          string
-}
+	DefaultConfig
 
-// UpdateAWSProfile updates or creates a single AWS profile to the AWS config file
-func (apc *AddProfileConfig) UpdateAWSProfile(iniFile *ini.File, profile *vault.ProfileSection, sourceProfile *string) error {
-	apc.Logger.Printf("Adding the profile %q to the AWS config file", profile.Name)
-	sectionName := fmt.Sprintf("profile %s", profile.Name)
-
-	// Get or create section before updating
-	var err error
-	var section *ini.Section
-	section = iniFile.Section(sectionName)
-	if section == nil {
-		section, err = iniFile.NewSection(sectionName)
-		if err != nil {
-			return fmt.Errorf("error creating section %q: %w", profile.Name, err)
-		}
-	}
+	BaseProfileName string
+	AWSProfileName  string
+}
 
-	// Add the source profile when provided
-	if sourceProfile != nil {
-		_, err = section.NewKey("source_profile", *sourceProfile)
-		if err != nil {
-			return fmt.Errorf("unable to add source profile: %w", err)
-		}
+// Run orchestrates the tasks to add a new profile
+func (apc *AddProfileConfig) Run() error {
+	if err := apc.AddProfile(); err != nil {
+		return err
 	}
 
-	if err = section.ReflectFrom(&profile); err != nil {
-		return fmt.Errorf("error mapping profile to ini file: %w", err)
-	}
-	_, err = section.NewKey("output", apc.Output)
-	if err != nil {
-		return fmt.Errorf("unable to add output key: %w", err)
-	}
 	return nil
 }
 
@@ -212,9 +177,6 @@ func addProfileFunction(cmd *cobra.Command, args []string) error {
 	iamRole := v.GetString(IAMRoleFlag)
 	output := v.GetString(OutputFlag)
 
-	// Validator used to validate input options for MFA
-	validate = validator.New()
-
 	// initialize things
 	partition, err := getPartition(awsRegion)
 	if err != nil {
@@ -226,23 +188,24 @@ func addProfileFunction(cmd *cobra.Command, args []string) error {
 		logger.Fatal(err)
 	}
 
-	addProfileConfig := AddProfileConfig{
-		// Config
-		Logger: logger,
-		Config: config,
+	// Setup new config
+	addProfileConfig := AddProfileConfig{}
 
-		// Profile Inputs
-		IAMRole:   iamRole,
-		Region:    awsRegion,
-		Partition: partition,
-		Output:    output,
+	// Config
+	addProfileConfig.Logger = logger
+	addProfileConfig.Config = config
 
-		// Profiles
-		AWSProfileAccounts: awsProfileAccount,
-		AWSProfileName:     awsProfile,
-	}
+	// Profile Inputs
+	addProfileConfig.IAMRole = iamRole
+	addProfileConfig.Region = awsRegion
+	addProfileConfig.Partition = partition
+	addProfileConfig.Output = output
+
+	// Profiles
+	addProfileConfig.AWSProfileAccounts = awsProfileAccount
+	addProfileConfig.AWSProfileName = awsProfile
 
-	if err := addProfileConfig.AddProfile(); err != nil {
+	if err := addProfileConfig.Run(); err != nil {
 		logger.Fatal(err)
 	}
 

cmd/add_profile_test.go

@@ -62,20 +62,25 @@ output=json
 
 	mfaSerial := "arn:aws:iam::111111111111:mfa/test-user"
 	addProfileConfig := AddProfileConfig{
-		// Config
-		Logger: suite.logger,
-		Config: config,
-
-		// Profile Inputs
-		IAMRole:   "test-role",
-		Region:    "us-west-2",
-		Partition: "aws",
-		Output:    "json",
-
-		// Profiles
-		AWSProfileAccounts: []string{"test-id-new:123456789012"},
-		AWSProfileName:     "test-id",
+		AWSProfileName: "test-id",
 	}
+
+	// Config
+	addProfileConfig.Logger = suite.logger
+	addProfileConfig.Config = config
+
+	// Profile Inputs
+	addProfileConfig.IAMUser = ""
+	addProfileConfig.IAMRole = "test-role"
+	addProfileConfig.Region = "us-west-2"
+	addProfileConfig.Partition = "aws"
+	addProfileConfig.Output = "json"
+
+	// Profiles
+	addProfileConfig.AWSProfileAccounts = []string{"test-id-new:123456789012"}
+	addProfileConfig.MFASerial = mfaSerial
+
+	// Add the profile
 	err = addProfileConfig.AddProfile()
 	suite.NoError(err)
 

cmd/common.go

@@ -14,6 +14,8 @@ import (
 	"gopkg.in/ini.v1"
 )
 
+const maxMFATokenPromptAttempts = 5
+
 func promptMFAtoken(messagePrefix string, logger *log.Logger) string {
 	var token string
 	for attempts := maxMFATokenPromptAttempts; token == "" && attempts > 0; attempts-- {
@@ -53,9 +55,25 @@ func getMFATokenPair(logger *log.Logger) MFATokenPair {
 	return mfaTokenPair
 }
 
+// DefaultConfig is the standard config struct for managing subcommand input
+type DefaultConfig struct {
+	Logger *log.Logger
+	Config *vault.ConfigFile
+
+	IAMUser   string
+	IAMRole   string
+	Partition string
+	Region    string
+	Output    string
+
+	AWSProfileAccounts []string
+	AWSProfiles        []vault.ProfileSection
+	MFASerial          string
+}
+
 // UpdateAWSProfile updates or creates a single AWS profile to the AWS config file
-func UpdateAWSProfile(iniFile *ini.File, profile *vault.ProfileSection, sourceProfile *string, output string, logger *log.Logger) error {
-	logger.Printf("Adding the profile %q to the AWS config file", profile.Name)
+func (dc *DefaultConfig) UpdateAWSProfile(iniFile *ini.File, profile *vault.ProfileSection, sourceProfile *string) error {
+	dc.Logger.Printf("Adding the profile %q to the AWS config file", profile.Name)
 
 	sectionName := fmt.Sprintf("profile %s", profile.Name)
 
@@ -81,7 +99,7 @@ func UpdateAWSProfile(iniFile *ini.File, profile *vault.ProfileSection, sourcePr
 	if err = section.ReflectFrom(&profile); err != nil {
 		return fmt.Errorf("error mapping profile to ini file: %w", err)
 	}
-	_, err = section.NewKey("output", output)
+	_, err = section.NewKey("output", dc.Output)
 	if err != nil {
 		return fmt.Errorf("unable to add output key: %w", err)
 	}

cmd/setup.go

@@ -25,7 +25,6 @@ import (
 )
 
 const maxNumAccessKeys = 2
-const maxMFATokenPromptAttempts = 5
 
 var validate *validator.Validate
 
@@ -88,73 +87,64 @@ func SetupUserCheckConfig(v *viper.Viper) error {
 
 // SetupConfig holds information for the AWS user being configured by this script
 type SetupConfig struct {
-	Logger     *log.Logger
-	Config     *vault.ConfigFile
+	DefaultConfig
+
 	QrTempFile *os.File
 	Keyring    *keyring.Keyring
 	NoMFA      bool
 
-	IAMUser   string
-	IAMRole   string
-	Partition string
-	Region    string
-	Output    string
-
-	BaseProfileName    string
-	BaseProfile        *vault.ProfileSection
-	AWSProfileAccounts []string
-	AWSProfiles        []vault.ProfileSection
-	MFASerial          string
+	BaseProfileName string
+	BaseProfile     *vault.ProfileSection
 
 	AccessKeyID     string
 	SecretAccessKey string
 }
 
-// Setup orchestrates the tasks to create the user's MFA and rotate access
-// keys.
-func (sc *SetupConfig) Setup() {
+// Run orchestrates the tasks to create the user's MFA and rotate access keys.
+func (sc *SetupConfig) Run() error {
 	err := sc.PromptAccessCredentials()
 	if err != nil {
-		sc.Logger.Fatal(err)
+		return err
 	}
 
 	err = sc.AddVaultProfile()
 	if err != nil {
-		sc.Logger.Fatal(err)
+		return err
 	}
 
 	if sc.NoMFA {
 		err = sc.GetMFADevice()
 		if err != nil {
-			sc.Logger.Fatal(err)
+			return err
 		}
 	} else {
 		err = sc.CreateVirtualMFADevice()
 		if err != nil {
-			sc.Logger.Fatal(err)
+			return err
 		}
 
 		err = sc.EnableVirtualMFADevice()
 		if err != nil {
-			sc.Logger.Fatal(err)
+			return err
 		}
 	}
 
 	err = sc.UpdateAWSConfigFile()
 	if err != nil {
-		sc.Logger.Fatal(err)
+		return err
 	}
 
 	err = sc.RemoveVaultSession()
 	if err != nil {
-		sc.Logger.Fatal(err)
+		return err
 	}
 
 	err = sc.RotateAccessKeys()
 	if err != nil {
-		sc.Logger.Fatal(err)
+		return err
 	}
 
+	return nil
 }
 
 // PromptAccessCredentials prompts the user for their AWS access key ID and
@@ -401,40 +391,6 @@ func (sc *SetupConfig) AddVaultProfile() error {
 	return nil
 }
 
-// UpdateAWSProfile updates or creates a single AWS profile to the AWS config file
-func (sc *SetupConfig) UpdateAWSProfile(iniFile *ini.File, profile *vault.ProfileSection, sourceProfile *string) error {
-	sc.Logger.Printf("Adding the profile %q to the AWS config file", profile.Name)
-	sectionName := fmt.Sprintf("profile %s", profile.Name)
-
-	// Get or create section before updating
-	var err error
-	var section *ini.Section
-	section = iniFile.Section(sectionName)
-	if section == nil {
-		section, err = iniFile.NewSection(sectionName)
-		if err != nil {
-			return fmt.Errorf("error creating section %q: %w", profile.Name, err)
-		}
-	}
-
-	// Add the source profile when provided
-	if sourceProfile != nil {
-		_, err = section.NewKey("source_profile", *sourceProfile)
-		if err != nil {
-			return fmt.Errorf("unable to add source profile: %w", err)
-		}
-	}
-
-	if err = section.ReflectFrom(&profile); err != nil {
-		return fmt.Errorf("error mapping profile to ini file: %w", err)
-	}
-	_, err = section.NewKey("output", sc.Output)
-	if err != nil {
-		return fmt.Errorf("unable to add output key: %w", err)
-	}
-	return nil
-}
-
 // UpdateAWSConfigFile adds the user's AWS profile to the AWS config file
 func (sc *SetupConfig) UpdateAWSConfigFile() error {
 	sc.Logger.Printf("Updating the AWS config file: %s", sc.Config.Path)
@@ -588,25 +544,31 @@ func setupUserFunction(cmd *cobra.Command, args []string) error {
 
 	setupConfig := SetupConfig{
 		// Config
-		Logger:     logger,
-		Config:     config,
 		QrTempFile: tempfile,
 		Keyring:    keyring,
 		NoMFA:      noMFA,
 
-		// Profile Inputs
-		IAMUser:   iamUser,
-		IAMRole:   iamRole,
-		Region:    awsRegion,
-		Partition: partition,
-		Output:    output,
-
 		// Profiles
-		BaseProfileName:    baseProfileName,
-		AWSProfileAccounts: awsProfileAccount,
+		BaseProfileName: baseProfileName,
 	}
 
-	setupConfig.Setup()
+	// Config
+	setupConfig.Logger = logger
+	setupConfig.Config = config
+
+	// Profile Inputs
+	setupConfig.IAMUser = iamUser
+	setupConfig.IAMRole = iamRole
+	setupConfig.Region = awsRegion
+	setupConfig.Partition = partition
+	setupConfig.Output = output
+
+	// Profiles
+	setupConfig.AWSProfileAccounts = awsProfileAccount
+
+	if err := setupConfig.Run(); err != nil {
+		logger.Fatal(err)
+	}
 
 	// If we got this far, we win
 	logger.Println("Victory!")