Fixed assume role without MFA

ivanfoo · Oct 9, 2018 · 2bb254e · 2bb254e
1 parent add93f3
commit 2bb254e
Showing 1 changed file with 19 additions and 7 deletions.
diff --git a/commands/assume.go b/commands/assume.go
@@ -107,13 +107,25 @@ func (c *AssumeCommand) askForTokenCode() string {
 
 func (c *AssumeCommand) getCredentials(serialNumber string, tokenCode string, roleArn string) (*awsCredentials, error) {
 	service := sts.New(c.session)
-	role, err := service.AssumeRole(&sts.AssumeRoleInput{
-		RoleSessionName: aws.String(c.Alias),
-		DurationSeconds: aws.Int64(c.Duration),
-		RoleArn:         aws.String(roleArn),
-		SerialNumber:    aws.String(serialNumber),
-		TokenCode:       aws.String(tokenCode),
-	})
+
+	var stsRequest *sts.AssumeRoleInput
+	if c.MFA {
+		stsRequest = &sts.AssumeRoleInput{
+			RoleSessionName: aws.String(c.Alias),
+			DurationSeconds: aws.Int64(c.Duration),
+			RoleArn:         aws.String(roleArn),
+			SerialNumber:    aws.String(serialNumber),
+			TokenCode:       aws.String(tokenCode),
+		}
+	} else {
+		stsRequest = &sts.AssumeRoleInput{
+			RoleSessionName: aws.String(c.Alias),
+			DurationSeconds: aws.Int64(c.Duration),
+			RoleArn:         aws.String(roleArn),
+		}
+	}
+
+	role, err := service.AssumeRole(stsRequest)
 
 	if err != nil {
 		return nil, err