Merge pull request #1 from dekachiri/feature/support_switch_role

Switch Role環境で利用可能にする

Author: akeboshi

README.md

@@ -4,10 +4,44 @@ aws cli がない状況でも RDS IAM Token を取得したく生み出された
 
 ## 使い方
 
-```
+### 通常のIAMユーザ + MFA認証の場合
+
+```bash
 export AWS_ACCESS_KEY_ID="FOO"
-export AWS_SECRET_ACCESS_KE="BAR"
+export AWS_SECRET_ACCESS_KEY="BAR"
 export AWS_DEFAULT_REGION="ap-northeast-1"
 
 ./RDSIamToken -user ${RDSのユーザ名} -endpoint ${RDSのエンドポイント}:${RDSのポート} -mfaARN ${arn:aws:iam::1234567890:mfa/xxx}
 ```
+
+### Switch Role（AssumeRole）を使用する場合
+
+```bash
+export AWS_ACCESS_KEY_ID="FOO"
+export AWS_SECRET_ACCESS_KEY="BAR"
+export AWS_DEFAULT_REGION="ap-northeast-1"
+
+./RDSIamToken -user ${RDSのユーザ名} -endpoint ${RDSのエンドポイント}:${RDSのポート} -mfaARN ${arn:aws:iam::1234567890:mfa/xxx} -roleArn ${arn:aws:iam::1234567890:role/MyRole}
+```
+
+### カスタムセッション名を指定する場合
+
+```bash
+./RDSIamToken -user ${RDSのユーザ名} -endpoint ${RDSのエンドポイント}:${RDSのポート} -mfaARN ${arn:aws:iam::1234567890:mfa/xxx} -roleArn ${arn:aws:iam::1234567890:role/MyRole} -roleSessionName "CustomSessionName"
+```
+
+### プロファイルを使用する場合
+
+```bash
+./RDSIamToken -profile ${プロファイル名} -user ${RDSのユーザ名} -endpoint ${RDSのエンドポイント}:${RDSのポート} -mfaARN ${arn:aws:iam::1234567890:mfa/xxx} -roleArn ${arn:aws:iam::1234567890:role/MyRole}
+```
+
+### パラメータ説明
+
+- `-user`: DBのユーザー名（例: iam_operator）
+- `-endpoint`: DBのエンドポイント（例: mydb.rds.amazonaws.com:3306）
+- `-mfaARN`: MFAデバイスのARN（例: arn:aws:iam::1234567890:mfa/username）
+- `-roleArn`: Switch Roleで使用するロールのARN（オプション、例: arn:aws:iam::1234567890:role/MyRole）
+- `-roleSessionName`: AssumeRoleで使用するセッション名（オプション、デフォルト: RDSIAMTokenSession）
+- `-profile`: 使用するAWSプロファイル名（オプション）
+- `-region`: 使用するAWSリージョン（デフォルト: ap-northeast-1）

main.go

@@ -14,7 +14,7 @@ import (
 	"flag"
 )
 
-func buildRDSTokenWithMFA(profile, defaultRegion, mfaARN, endpoint, dbUser string) (string, error) {
+func buildRDSTokenWithMFA(profile, defaultRegion, mfaARN, roleArn, roleSessionName, endpoint, dbUser string) (string, error) {
 	// Profile の読み込み
 	cfg, err := config.LoadDefaultConfig(context.Background(),
 		config.WithSharedConfigProfile(profile),
@@ -29,20 +29,47 @@ func buildRDSTokenWithMFA(profile, defaultRegion, mfaARN, endpoint, dbUser strin
 	fmt.Fprintln(os.Stderr, "Please input MFA code:")
     fmt.Scan(&mfaToken)
 	stsClient := sts.NewFromConfig(cfg)
-	creds, err := stsClient.GetSessionToken(context.Background(), &sts.GetSessionTokenInput{
-		TokenCode:       aws.String(mfaToken),
-		SerialNumber:    aws.String(mfaARN),
-		DurationSeconds: aws.Int32(3600),
-	})
 
-	if err != nil {
-		return "", fmt.Errorf("Failed to get session token: %w", err)
+	var finalCredentials aws.Credentials
+
+	if roleArn != "" {
+		// Switch Role を使用する場合
+		assumeRoleOutput, err := stsClient.AssumeRole(context.Background(), &sts.AssumeRoleInput{
+			RoleArn:         aws.String(roleArn),
+			RoleSessionName: aws.String(roleSessionName),
+			TokenCode:       aws.String(mfaToken),
+			SerialNumber:    aws.String(mfaARN),
+			DurationSeconds: aws.Int32(3600),
+		})
+		if err != nil {
+			return "", fmt.Errorf("Failed to assume role: %w", err)
+		}
+		finalCredentials = aws.Credentials{
+			AccessKeyID:     *assumeRoleOutput.Credentials.AccessKeyId,
+			SecretAccessKey: *assumeRoleOutput.Credentials.SecretAccessKey,
+			SessionToken:    *assumeRoleOutput.Credentials.SessionToken,
+		}
+	} else {
+		// 通常のMFA認証の場合
+		creds, err := stsClient.GetSessionToken(context.Background(), &sts.GetSessionTokenInput{
+			TokenCode:       aws.String(mfaToken),
+			SerialNumber:    aws.String(mfaARN),
+			DurationSeconds: aws.Int32(3600),
+		})
+		if err != nil {
+			return "", fmt.Errorf("Failed to get session token: %w", err)
+		}
+		finalCredentials = aws.Credentials{
+			AccessKeyID:     *creds.Credentials.AccessKeyId,
+			SecretAccessKey: *creds.Credentials.SecretAccessKey,
+			SessionToken:    *creds.Credentials.SessionToken,
+		}
 	}
 
 	cfg.Credentials = aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(
-		*creds.Credentials.AccessKeyId,
-		*creds.Credentials.SecretAccessKey,
-		*creds.Credentials.SessionToken,
+		finalCredentials.AccessKeyID,
+		finalCredentials.SecretAccessKey,
+		finalCredentials.SessionToken,
 	))
 
 	// RDS の認証情報の取得
@@ -60,6 +87,8 @@ func main() {
 	// 使用するAWSプロファイル名を設定
 	profile := flag.String("profile", "", "AWS の Profile 名")
 	mfaARN := flag.String("mfaARN", "", "mfaARN 例) arn:aws:iam::xxxxxxxx:mfa/isao-aruga")
+	roleArn := flag.String("roleArn", "", "Switch Roleで使用するロールのARN 例) arn:aws:iam::xxxxxxxx:role/MyRole")
+	roleSessionName := flag.String("roleSessionName", "RDSIAMTokenSession", "AssumeRoleで使用するセッション名")
 	dbUser := flag.String("user", "", "DBのユーザー名 例) iam_operator")
 	endpoint := flag.String("endpoint", "", "DBのエンドポイント 例) sms-main-xxxxxx.rds.amazonaws.com:3306")
 	defaultRegion := flag.String("region", "ap-northeast-1", "利用するAWSのリージョン")
@@ -80,7 +109,7 @@ func main() {
 		os.Exit(1)
 	}
 
-	token, err := buildRDSTokenWithMFA(*profile, *defaultRegion, *mfaARN, *endpoint, *dbUser)
+	token, err := buildRDSTokenWithMFA(*profile, *defaultRegion, *mfaARN, *roleArn, *roleSessionName, *endpoint, *dbUser)
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(1)