Merge pull request #17 from uc4w6c/feature/set-session-name

feat: Set session name when assuming role

Author: uc4w6c

gradle.properties

@@ -4,7 +4,7 @@ pluginGroup = com.github.uc4w6c.bedrockassistant
 pluginName = BedrockAssistant
 pluginRepositoryUrl = https://github.com/uc4w6c/BedrockAssistant
 # SemVer format -> https://semver.org
-pluginVersion = 0.1.2
+pluginVersion = 0.1.3
 
 # Supported build number ranges and IntelliJ Platform versions -> https://plugins.jetbrains.com/docs/intellij/build-number-ranges.html
 pluginSinceBuild = 242

src/main/java/com/github/uc4w6c/bedrockassistant/dao/AssumeRoleDao.java

@@ -15,13 +15,15 @@ public class AssumeRoleDao {
   public AssumeRoleResponse getToken(AssumeRoleRequestEntity requestEntity) {
     AssumeRoleRequest.Builder requestBuilder = AssumeRoleRequest
         .builder()
-        .roleArn(requestEntity.roleArn())
-        .roleSessionName(UUID.randomUUID().toString());
+        .roleArn(requestEntity.roleArn());
 
     if (requestEntity.mfaSerial().isPresent()) {
       requestBuilder.serialNumber(requestEntity.mfaSerial().get());
       requestBuilder.tokenCode(requestEntity.tokenCode().get());
     }
+    requestBuilder.roleSessionName(
+        requestEntity.roleSessionName()
+            .orElse(UUID.randomUUID().toString()));
     AssumeRoleRequest request = requestBuilder.build();
 
     StsClientBuilder stsClientBuilder = StsClient.builder();

src/main/java/com/github/uc4w6c/bedrockassistant/dao/entity/AssumeRoleRequestEntity.java

@@ -6,13 +6,15 @@ public record AssumeRoleRequestEntity(
     String roleArn,
     Optional<String> assumeRoleRegion,
     Optional<String> mfaSerial,
-    Optional<String> tokenCode
+    Optional<String> tokenCode,
+    Optional<String> roleSessionName
 ) {
   public static class Builder {
     String roleArn = null;
     Optional<String> assumeRoleRegion = Optional.empty();
     Optional<String> mfaSerial = Optional.empty();
     Optional<String> tokenCode = Optional.empty();
+    Optional<String> roleSessionName = Optional.empty();
 
     public Builder roleArn(String roleArn) {
       this.roleArn = roleArn;
@@ -34,6 +36,11 @@ public Builder tokenCode(String tokenCode) {
       return this;
     }
 
+    public Builder roleSessionName(String roleSessionName) {
+      this.roleSessionName = Optional.of(roleSessionName);
+      return this;
+    }
+
     public AssumeRoleRequestEntity build() {
       if (this.roleArn == null) throw new IllegalArgumentException();
       if (this.mfaSerial.isPresent() != this.tokenCode.isPresent()) throw new IllegalArgumentException();
@@ -42,7 +49,8 @@ public AssumeRoleRequestEntity build() {
           this.roleArn,
           this.assumeRoleRegion,
           this.mfaSerial,
-          this.tokenCode);
+          this.tokenCode,
+          this.roleSessionName);
     }
   }
 }

src/main/java/com/github/uc4w6c/bedrockassistant/domain/AwsProfile.java

@@ -9,12 +9,14 @@
  * @param region
  * @param roleArn
  * @param mfaSerial
+ * @param roleSessionName
  */
 public record AwsProfile(
     String name,
     Optional<String> region,
     Optional<String> roleArn,
-    Optional<String> mfaSerial
+    Optional<String> mfaSerial,
+    Optional<String> roleSessionName
 ) {
   @Override
   public boolean equals(Object o) {

src/main/java/com/github/uc4w6c/bedrockassistant/repository/AssumeRoleRepository.java

@@ -27,6 +27,9 @@ public AwsCredentials getTokenWithMfa(AwsProfile profile, String tokenCode) {
     if (profile.region().isPresent()) {
       builder.assumeRoleRegion(profile.region().get());
     }
+    if (profile.roleSessionName().isPresent()) {
+      builder.roleSessionName(profile.roleSessionName().get());
+    }
     builder.mfaSerial(profile.mfaSerial().get());
     builder.tokenCode(tokenCode);
 
@@ -39,11 +42,16 @@ public AwsCredentials getTokenWithMfa(AwsProfile profile, String tokenCode) {
   }
 
   public AwsCredentials getToken(AwsProfile profile) {
-    AssumeRoleRequestEntity request = new AssumeRoleRequestEntity.Builder()
-        .roleArn(profile.roleArn().get())
-        .build();
+    AssumeRoleRequestEntity.Builder builder = new AssumeRoleRequestEntity.Builder();
+    builder.roleArn(profile.roleArn().get());
+    if (profile.region().isPresent()) {
+      builder.assumeRoleRegion(profile.region().get());
+    }
+    if (profile.roleSessionName().isPresent()) {
+      builder.roleSessionName(profile.roleSessionName().get());
+    }
 
-    AssumeRoleResponse response = assumeRoleDao.getToken(request);
+    AssumeRoleResponse response = assumeRoleDao.getToken(builder.build());
     return new AwsCredentials(
         response.credentials().accessKeyId(),
         response.credentials().secretAccessKey(),

src/main/java/com/github/uc4w6c/bedrockassistant/repository/ProfileRepository.java

@@ -23,8 +23,9 @@ public Set<AwsProfile> getProfiles() {
             new AwsProfile(
                 entry.getKey(),
                 entry.getValue().property("region"),
-                entry.getValue().property("role-arn"),
-                entry.getValue().property("mfa-serial")))
+                entry.getValue().property("role_arn"),
+                entry.getValue().property("mfa_serial"),
+                entry.getValue().property("role_session_name")))
         .collect(Collectors.toSet());
   }
 
@@ -33,6 +34,7 @@ public Optional<AwsProfile> getProfile(String profileName) {
     return optionalProfile.map(profile -> new AwsProfile(profile.name(),
         profile.property("region"),
         profile.property("role_arn"),
-        profile.property("mfa_serial")));
+        profile.property("mfa_serial"),
+        profile.property("role_session_name")));
   }
 }