/*
 * Copyright 2016 Netflix, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.netflix.spinnaker.clouddriver.aws.security;

import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.auth.AWSSessionCredentials;
import com.amazonaws.auth.AWSSessionCredentialsProvider;
import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider;
import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClient;
import com.netflix.spinnaker.clouddriver.aws.security.sdkclient.SpinnakerAwsRegionProvider;
import java.io.Closeable;

public class NetflixSTSAssumeRoleSessionCredentialsProvider
    implements AWSSessionCredentialsProvider, Closeable {

  private final String accountId;
  private final STSAssumeRoleSessionCredentialsProvider delegate;

  public NetflixSTSAssumeRoleSessionCredentialsProvider(
      AWSCredentialsProvider longLivedCredentialsProvider,
      String roleArn,
      String roleSessionName,
      String accountId,
      String externalId) {
    this.accountId = accountId;

    var chain = new SpinnakerAwsRegionProvider();
    var region = chain.getRegion();

    delegate =
        new STSAssumeRoleSessionCredentialsProvider.Builder(roleArn, roleSessionName)
            .withExternalId(externalId)
            .withStsClient(
                AWSSecurityTokenServiceClient.builder()
                    .withCredentials(longLivedCredentialsProvider)
                    .withRegion(region)
                    .build())
            .build();

    /**
     * Need to explicitly set sts region if GovCloud or China as per
     * https://docs.aws.amazon.com/AWSJavaSDK/latest/javadoc/com/amazonaws/auth/STSAssumeRoleSessionCredentialsProvider.html
     */
    if (roleArn.contains("aws-us-gov")) {
      delegate.setSTSClientEndpoint("sts.us-gov-west-1.amazonaws.com");
    }
    if (roleArn.contains("aws-cn")) {
      delegate.setSTSClientEndpoint("sts.cn-north-1.amazonaws.com.cn");
    }
  }

  public String getAccountId() {
    return accountId;
  }

  @Override
  public AWSSessionCredentials getCredentials() {
    return delegate.getCredentials();
  }

  @Override
  public void refresh() {
    delegate.refresh();
  }

  @Override
  public void close() {
    delegate.close();
  }
}

@@ -862,32 +862,108 @@
 ntialsProvider;%0A
+import com.amazonaws.client.builder.AwsClientBuilder.EndpointConfiguration;%0A
 import com.amazo
@@ -1707,181 +1707,31 @@
     
-delegate =%0A        new STSAssumeRoleSessionCredentialsProvider.Builder(roleArn, roleSessionName)%0A            .withExternalId(externalId)%0A            .withStsClient(%0A        
+var stsClientBuilder =%0A
     
@@ -1773,37 +1773,16 @@
 uilder()
-%0A                    
 .withCre
@@ -1823,315 +1823,194 @@
 der)
+;%0A
 %0A    
-                .withRegion(region)%0A                    .build())%0A            .build();%0A
+if (roleArn.contains(%22aws-us-gov%22)) %7B%0A      stsClientBuilder.withEndpointConfiguration(
 %0A    
-/**%0A
      
-* Need to explicitly set sts region if GovCloud or China as per%0A     * https://docs.aws.amazon.com/AWSJavaSDK/latest/javadoc/com/amazonaws/auth/STSAssumeRoleSessionCredentialsProvider.html%0A     */%0A   
+ new EndpointConfiguration(%22sts.us-gov-west-1.amazonaws.com%22, region));%0A    %7D else
  if 
@@ -2024,38 +2024,34 @@
 n.contains(%22aws-
-us-gov
+cn
 %22)) %7B%0A      dele
@@ -2050,195 +2050,411 @@
     
-delegate.setSTSClientEndpoint(%22sts.us-gov-west-1.amazonaws.com%22);%0A    %7D%0A    if (roleArn.contains(%22aws-cn%22)) %7B%0A      delegate.setSTSClientEndpoi
+stsClientBuilder.withEndpointConfiguration(%0A          new EndpointConfiguration(%22sts.cn-north-1.amazonaws.com.cn%22, region));%0A    %7D else %7B%0A      stsClientBuilder.withRegion(region);%0A    %7D%0A%0A    delegate =%0A        new STSAssumeRoleSessionCredentialsProvider.Builder(roleArn, roleSessionName)%0A            .withExternalId(externalId)%0A            .withStsClie
 nt(
-%22
 sts
-.cn-north-1.amazonaws.com.cn%22);%0A    %7D
+ClientBuilder.build())%0A            .build();
 %0A  %7D

Fix STS role provider for AWS GovCloud and AWS CN (#5202)