com.amazonaws.services.securitytoken.model.AssumeRoleResult Java Examples

/**
 * Gets the credentials.
 *
 * @param account the account
 * @param roleName the role name
 * @return the credentials
 */
public  BasicSessionCredentials getCredentials(String account,String roleName){
	
	BasicSessionCredentials baseAccntCreds = getBaseAccountCredentials(baseAccount,baseRegion,roleName);
	if(baseAccount.equals(account)){
		return baseAccntCreds;
	}
	AWSSecurityTokenServiceClientBuilder stsBuilder = AWSSecurityTokenServiceClientBuilder.standard().withCredentials( new AWSStaticCredentialsProvider(baseAccntCreds)).withRegion(baseRegion);
	AWSSecurityTokenService stsClient = stsBuilder.build();
    AssumeRoleRequest assumeRequest = new AssumeRoleRequest().withRoleArn(getRoleArn(account,roleName)).withRoleSessionName("pic-ro-"+account);
    AssumeRoleResult assumeResult = stsClient.assumeRole(assumeRequest);
    return  new BasicSessionCredentials(
            assumeResult.getCredentials()
                        .getAccessKeyId(), assumeResult.getCredentials().getSecretAccessKey(),
            assumeResult.getCredentials().getSessionToken());
}
 

/**
 * Gets the temp credentials using cred provider.
 *
 * @param roleArnWithAdequateAccess
 *            the role arn with adequate access
 * @param region
 *            the region
 * @param acp
 *            the acp
 * @param validForSeconds
 *            the valid for seconds
 * @return the temp credentials using cred provider
 */
private BasicSessionCredentials getTempCredentialsUsingCredProvider(String roleArnWithAdequateAccess,
        Regions region, AWSCredentialsProvider acp, Integer validForSeconds) {
    if (null == region) { // cloud trail case
        region = Regions.DEFAULT_REGION;
    }
    AWSSecurityTokenServiceClientBuilder stsBuilder = AWSSecurityTokenServiceClientBuilder.standard()
            .withCredentials(acp).withRegion(region);
    AWSSecurityTokenService sts = stsBuilder.build();
    AssumeRoleRequest assumeRequest = new AssumeRoleRequest().withRoleArn(roleArnWithAdequateAccess)
            .withDurationSeconds(validForSeconds).withRoleSessionName(PacmanSdkConstants.DEFAULT_SESSION_NAME);
    logger.debug("assume role request " + assumeRequest.toString());
    AssumeRoleResult assumeResult = sts.assumeRole(assumeRequest);
    logger.debug("assume role response " + assumeResult.toString());
    BasicSessionCredentials temporaryCredentials = new BasicSessionCredentials(assumeResult.getCredentials()
            .getAccessKeyId(), assumeResult.getCredentials().getSecretAccessKey(), assumeResult.getCredentials()
            .getSessionToken());

    return temporaryCredentials;
}
 

private BasicSessionCredentials getFreshCredentials(AWSEnvironment environment) throws GatekeeperException{

        logger.info("Assuming role for environment " + environment.getAccount() + " on region " + environment.getRegion()
                + " with timeout of " + (gatekeeperAwsProperties.getSessionTimeout() / 1000) + " seconds (with " + (gatekeeperAwsProperties.getSessionTimeoutPad() / 1000) + " padding.)");

        AssumeRoleRequest assumeRequest = new AssumeRoleRequest()
                .withRoleArn(getRoleArn(environment.getAccount()))
                .withDurationSeconds((gatekeeperAwsProperties.getSessionTimeout() + gatekeeperAwsProperties.getSessionTimeoutPad()) / 1000)
                .withRoleSessionName("GATEKEEPER_APP");

        AssumeRoleResult assumeResult = awsSecurityTokenServiceClient.assumeRole(assumeRequest);

        return new BasicSessionCredentials(
                assumeResult.getCredentials().getAccessKeyId(),
                assumeResult.getCredentials().getSecretAccessKey(),
                assumeResult.getCredentials().getSessionToken());

    }
 

private BasicSessionCredentials getFreshCredentials(AWSEnvironment environment) throws GatekeeperException{

        logger.info("Assuming role for environment " + environment.getAccount() + " on region " + environment.getRegion()
                + " with timeout of " + (sessionTimeout / 1000) + " seconds (with " + (sessionTimeoutPad / 1000) + " padding.)");

        AssumeRoleRequest assumeRequest = new AssumeRoleRequest()
                .withRoleArn(getRoleArn(environment.getAccount()))
                .withDurationSeconds((sessionTimeout + sessionTimeoutPad) / 1000)
                .withRoleSessionName("GATEKEEPER_APP");

        AssumeRoleResult assumeResult = awsSecurityTokenServiceClient.assumeRole(assumeRequest);

        return new BasicSessionCredentials(
                assumeResult.getCredentials().getAccessKeyId(),
                assumeResult.getCredentials().getSecretAccessKey(),
                assumeResult.getCredentials().getSessionToken());

    }
 

@Override
public AssumeRoleResult assumeRole(AWSSecurityTokenServiceClient awsSecurityTokenServiceClient, AssumeRoleRequest assumeRoleRequest)
{
    assertNotNull(assumeRoleRequest);

    if (assumeRoleRequest.getPolicy() != null && assumeRoleRequest.getPolicy().equals(MockAwsOperationsHelper.AMAZON_THROTTLING_EXCEPTION))
    {
        AmazonServiceException throttlingException = new AmazonServiceException("test throttling exception");
        throttlingException.setErrorCode("ThrottlingException");

        throw throttlingException;
    }

    AssumeRoleResult assumeRoleResult = new AssumeRoleResult();

    assumeRoleResult.setCredentials(new Credentials(MOCK_AWS_ASSUMED_ROLE_ACCESS_KEY, MOCK_AWS_ASSUMED_ROLE_SECRET_KEY, MOCK_AWS_ASSUMED_ROLE_SESSION_TOKEN,
        new Date(System.currentTimeMillis() + 1000 * assumeRoleRequest.getDurationSeconds())));

    return assumeRoleResult;
}
 

public AwsSessionCredentials retrieveSessionCredentials(AwsCredentialView awsCredential) {
    String externalId = awsCredential.getExternalId();
    AssumeRoleRequest assumeRoleRequest = new AssumeRoleRequest()
            .withDurationSeconds(DEFAULT_SESSION_CREDENTIALS_DURATION)
            .withExternalId(StringUtils.isEmpty(externalId) ? deprecatedExternalId : externalId)
            .withRoleArn(awsCredential.getRoleArn())
            .withRoleSessionName(roleSessionName);
    LOGGER.debug("Trying to assume role with role arn {}", awsCredential.getRoleArn());
    try {
        AssumeRoleResult result = awsSecurityTokenServiceClient(awsCredential).assumeRole(assumeRoleRequest);
        Credentials credentialsResponse = result.getCredentials();

        String formattedExpirationDate = "";
        Date expirationTime = credentialsResponse.getExpiration();
        if (expirationTime != null) {
            formattedExpirationDate = new StdDateFormat().format(expirationTime);
        }
        LOGGER.debug("Assume role result credential: role arn: {}, expiration date: {}",
                awsCredential.getRoleArn(), formattedExpirationDate);

        return new AwsSessionCredentials(
                credentialsResponse.getAccessKeyId(),
                credentialsResponse.getSecretAccessKey(),
                credentialsResponse.getSessionToken(),
                credentialsResponse.getExpiration());
    } catch (SdkClientException e) {
        LOGGER.error("Unable to assume role. Check exception for details.", e);
        throw e;
    }
}
 

private static AWSCredentials getCredentials(String iamRole, String externalId) {
    if (isEmpty(iamRole)) return null;

    AWSSecurityTokenServiceClient sts = new AWSSecurityTokenServiceClient();

    int credsDuration = (int) (AWSCodeDeployPublisher.DEFAULT_TIMEOUT_SECONDS
                    * AWSCodeDeployPublisher.DEFAULT_POLLING_FREQUENCY_SECONDS);

    if (credsDuration > 3600) {
        credsDuration = 3600;
    }

    AssumeRoleResult assumeRoleResult = sts.assumeRole(new AssumeRoleRequest()
                    .withRoleArn(iamRole)
                    .withExternalId(externalId)
                    .withDurationSeconds(credsDuration)
                    .withRoleSessionName(AWSCodeDeployPublisher.ROLE_SESSION_NAME)
    );

    Credentials stsCredentials = assumeRoleResult.getCredentials();
    BasicSessionCredentials credentials = new BasicSessionCredentials(
            stsCredentials.getAccessKeyId(),
            stsCredentials.getSecretAccessKey(),
            stsCredentials.getSessionToken()
    );

    return credentials;
}
 

/**
 * Returns a set of temporary security credentials (consisting of an access key ID, a secret access key, and a security token) that can be used to access
 * the specified AWS resource.
 *
 * @param sessionName the session name that will be associated with the temporary credentials. The session name must be the same for an initial set of
 * credentials and an extended set of credentials if credentials are to be refreshed. The session name also is used to identify the user in AWS logs so it
 * should be something unique and useful to identify the caller/use.
 * @param awsRoleArn the AWS ARN for the role required to provide access to the specified AWS resource
 * @param awsRoleDurationSeconds the duration, in seconds, of the role session. The value can range from 900 seconds (15 minutes) to 3600 seconds (1 hour).
 * @param policy the temporary policy to apply to this request
 *
 * @return the assumed session credentials
 */
@Override
public Credentials getTemporarySecurityCredentials(AwsParamsDto awsParamsDto, String sessionName, String awsRoleArn, int awsRoleDurationSeconds,
    Policy policy)
{
    // Construct a new AWS security token service client using the specified client configuration to access Amazon S3.
    // A credentials provider chain will be used that searches for credentials in this order:
    // - Environment Variables - AWS_ACCESS_KEY_ID and AWS_SECRET_KEY
    // - Java System Properties - aws.accessKeyId and aws.secretKey
    // - Instance Profile Credentials - delivered through the Amazon EC2 metadata service

    ClientConfiguration clientConfiguration = new ClientConfiguration().withRetryPolicy(retryPolicyFactory.getRetryPolicy());

    // Only set the proxy hostname and/or port if they're configured.
    if (StringUtils.isNotBlank(awsParamsDto.getHttpProxyHost()))
    {
        clientConfiguration.setProxyHost(awsParamsDto.getHttpProxyHost());
    }
    if (awsParamsDto.getHttpProxyPort() != null)
    {
        clientConfiguration.setProxyPort(awsParamsDto.getHttpProxyPort());
    }

    AWSSecurityTokenServiceClient awsSecurityTokenServiceClient = new AWSSecurityTokenServiceClient(clientConfiguration);

    // Create the request.
    AssumeRoleRequest assumeRoleRequest = new AssumeRoleRequest();
    assumeRoleRequest.setRoleSessionName(sessionName);
    assumeRoleRequest.setRoleArn(awsRoleArn);
    assumeRoleRequest.setDurationSeconds(awsRoleDurationSeconds);
    if (policy != null)
    {
        assumeRoleRequest.setPolicy(policy.toJson());
    }

    // Get the temporary security credentials.
    AssumeRoleResult assumeRoleResult = stsOperations.assumeRole(awsSecurityTokenServiceClient, assumeRoleRequest);
    return assumeRoleResult.getCredentials();
}
 

@Test
public void testAssumeAWSRole() {
    MockCloudStore cloudStore = new MockCloudStore();
    cloudStore.awsEnabled = true;
    AssumeRoleResult mockResult = Mockito.mock(AssumeRoleResult.class);
    Credentials creds = Mockito.mock(Credentials.class);
    Mockito.when(creds.getAccessKeyId()).thenReturn("accesskeyid");
    Mockito.when(creds.getSecretAccessKey()).thenReturn("secretaccesskey");
    Mockito.when(creds.getSessionToken()).thenReturn("sessiontoken");
    Mockito.when(creds.getExpiration()).thenReturn(new Date());
    Mockito.when(mockResult.getCredentials()).thenReturn(creds);
    cloudStore.setAssumeRoleResult(mockResult);
    cloudStore.setReturnSuperAWSRole(true);

    AWSTemporaryCredentials awsCreds = cloudStore.assumeAWSRole("account", "syncer", "athenz.syncer", null, null);
    assertNotNull(awsCreds);
    assertEquals(awsCreds.getAccessKeyId(), "accesskeyid");
    assertEquals(awsCreds.getSessionToken(), "sessiontoken");
    assertEquals(awsCreds.getSecretAccessKey(), "secretaccesskey");
    cloudStore.close();
}
 

@Override
public void refresh() {
    if (!iamRoleArn.isEmpty()) {
        if (!haveCredentialsExpired()) {
            return;
        }

        AWSCredentialsProvider credentialsProvider = getBasicCredentialsOrDefaultChain(accessKey, secretKey);
        AWSCredentials credentials = credentialsProvider.getCredentials();

        AssumeRoleRequest assumeRequest = new AssumeRoleRequest()
                .withRoleArn(iamRoleArn)
                .withExternalId(externalId)
                .withDurationSeconds(3600)
                .withRoleSessionName(ROLE_SESSION_NAME);

        AssumeRoleResult assumeResult = new AWSSecurityTokenServiceClient(credentials).assumeRole(assumeRequest);

        roleCredentials = assumeResult.getCredentials();
    }
}
 

private AssumedRole assumeRole(final AWSSecurityTokenService sts) {
	final AssumeRoleRequest assumeRoleRequest = new AssumeRoleRequest().withRoleArn(this.roleArn)
					.withRoleSessionName(this.sessionName)
					.withDurationSeconds(this.durationInSeconds);
	Optional.ofNullable(this.externalId).ifPresent(assumeRoleRequest::setExternalId);
	Optional.ofNullable(this.policy).ifPresent(assumeRoleRequest::withPolicy);
	AssumeRoleResult assumeRoleResult = sts.assumeRole(assumeRoleRequest);
	return new AssumedRole(assumeRoleResult.getCredentials(), assumeRoleResult.getAssumedRoleUser());
}
 

/**
 * Gets the credentials.
 *
 * @param account the account
 * @param roleName the role name
 * @return the credentials
 */
public  BasicSessionCredentials getCredentials(String account,String roleName){
	BasicSessionCredentials baseAccntCreds = getBaseAccountCredentials(roleName);
	if(baseAccount.equals(account)){
		return baseAccntCreds;
	}
	AWSSecurityTokenServiceClientBuilder stsBuilder = AWSSecurityTokenServiceClientBuilder.standard().withCredentials( new AWSStaticCredentialsProvider(baseAccntCreds)).withRegion(baseRegion);
	AWSSecurityTokenService stsClient = stsBuilder.build();
    AssumeRoleRequest assumeRequest = new AssumeRoleRequest().withRoleArn(getRoleArn(account,roleName)).withRoleSessionName("pic-ro-"+account);
    AssumeRoleResult assumeResult = stsClient.assumeRole(assumeRequest);
    return  new BasicSessionCredentials(
            assumeResult.getCredentials()
                        .getAccessKeyId(), assumeResult.getCredentials().getSecretAccessKey(),
            assumeResult.getCredentials().getSessionToken());
}
 

@Test
public void testAssumeAWSRoleFailedCreds() {
    MockCloudStore cloudStore = new MockCloudStore();
    cloudStore.awsEnabled = true;
    AssumeRoleResult mockResult = Mockito.mock(AssumeRoleResult.class);
    Credentials creds = Mockito.mock(Credentials.class);
    Mockito.when(creds.getAccessKeyId()).thenReturn("accesskeyid");
    Mockito.when(creds.getSecretAccessKey()).thenReturn("secretaccesskey");
    Mockito.when(creds.getSessionToken()).thenReturn("sessiontoken");
    Mockito.when(creds.getExpiration()).thenReturn(new Date());
    Mockito.when(mockResult.getCredentials()).thenReturn(creds);
    cloudStore.setAssumeRoleResult(mockResult);
    cloudStore.setReturnSuperAWSRole(true);

    // add our key to the invalid cache

    cloudStore.putInvalidCacheCreds(cloudStore.getCacheKey("account", "syncer", "athenz.syncer", null, null));
    assertNull(cloudStore.assumeAWSRole("account", "syncer", "athenz.syncer", null, null));
    assertNull(cloudStore.assumeAWSRole("account", "syncer", "athenz.syncer", null, null));

    // now set the timeout to 1 second and sleep that long and after
    // that our test case should work as before

    cloudStore.invalidCacheTimeout = 1;
    try {
        Thread.sleep(1000);
    } catch (InterruptedException ignored) {
    }
    assertNotNull(cloudStore.assumeAWSRole("account", "syncer", "athenz.syncer", null, null));
    cloudStore.close();
}
 

@Override
public AssumeRoleResult assumeRole(AWSSecurityTokenServiceClient awsSecurityTokenServiceClient, AssumeRoleRequest assumeRoleRequest)
{
    return awsSecurityTokenServiceClient.assumeRole(assumeRoleRequest);
}
 

@Override
public Mono<Void> canAgentAssume(String iamRoleName) {
    return Mono.defer(() -> {
        long startTime = registry.clock().wallTime();

        // Check cache first
        Either<Boolean, Throwable> lastCheck = canAssumeCache.getIfPresent(iamRoleName);
        if (lastCheck != null) {
            return lastCheck.hasValue() ? Mono.empty() : Mono.error(lastCheck.getError());
        }

        // Must call AWS STS service
        return AwsReactorExt
                .<AssumeRoleRequest, AssumeRoleResult>toMono(
                        () -> new AssumeRoleRequest()
                                .withRoleSessionName("titusIamRoleValidation")
                                .withRoleArn(iamRoleName)
                                .withDurationSeconds(MIN_ASSUMED_ROLE_DURATION_SEC),
                        stsAgentClient::assumeRoleAsync
                )
                .flatMap(response -> {
                    logger.debug("Assumed into: {}", iamRoleName);
                    canAssumeCache.put(iamRoleName, Either.ofValue(true));
                    connectorMetrics.success(IamConnectorMetrics.IamMethods.CanAgentAssume, startTime);
                    return Mono.<Void>empty();
                })
                .onErrorMap(error -> {
                    logger.debug("Error: {}", error.getMessage());
                    connectorMetrics.failure(IamConnectorMetrics.IamMethods.CanAgentAssume, error, startTime);

                    String errorCode = ((AWSSecurityTokenServiceException) error).getErrorCode();
                    if ("AccessDenied".equals(errorCode)) {
                        // STS service returns access denied error with no additional clues. To get more insight we
                        // would have to make a call to IAM service, but this would require access to all client accounts.
                        IamConnectorException cannotAssumeError = IamConnectorException.iamRoleCannotAssume(iamRoleName, configuration.getDataPlaneAgentRoleArn());
                        canAssumeCache.put(iamRoleName, Either.ofError(cannotAssumeError));
                        return cannotAssumeError;
                    }
                    return IamConnectorException.iamRoleUnexpectedError(iamRoleName, error.getMessage());
                });
    });
}
 

@Test
public void testGetTemporarySecurityCredentials()
{
    // Create an AWS parameters DTO with proxy settings.
    AwsParamsDto awsParamsDto = new AwsParamsDto();
    awsParamsDto.setHttpProxyHost(HTTP_PROXY_HOST);
    awsParamsDto.setHttpProxyPort(HTTP_PROXY_PORT);

    // Specify the duration, in seconds, of the role session.
    int awsRoleDurationSeconds = INTEGER_VALUE;

    // Create an IAM policy.
    Policy policy = new Policy(STRING_VALUE);

    // Create a retry policy.
    RetryPolicy retryPolicy =
        new RetryPolicy(PredefinedRetryPolicies.DEFAULT_RETRY_CONDITION, PredefinedRetryPolicies.DEFAULT_BACKOFF_STRATEGY, INTEGER_VALUE, true);

    // Create the expected assume role request.
    AssumeRoleRequest assumeRoleRequest = new AssumeRoleRequest().withRoleArn(AWS_ROLE_ARN).withRoleSessionName(SESSION_NAME).withPolicy(policy.toJson())
        .withDurationSeconds(awsRoleDurationSeconds);

    // Create AWS credentials for API authentication.
    Credentials credentials = new Credentials();
    credentials.setAccessKeyId(AWS_ASSUMED_ROLE_ACCESS_KEY);
    credentials.setSecretAccessKey(AWS_ASSUMED_ROLE_SECRET_KEY);
    credentials.setSessionToken(AWS_ASSUMED_ROLE_SESSION_TOKEN);

    // Create an assume role result.
    AssumeRoleResult assumeRoleResult = new AssumeRoleResult();
    assumeRoleResult.setCredentials(credentials);

    // Mock the external calls.
    when(retryPolicyFactory.getRetryPolicy()).thenReturn(retryPolicy);
    when(stsOperations.assumeRole(any(AWSSecurityTokenServiceClient.class), eq(assumeRoleRequest))).thenReturn(assumeRoleResult);

    // Call the method under test.
    Credentials result = stsDaoImpl.getTemporarySecurityCredentials(awsParamsDto, SESSION_NAME, AWS_ROLE_ARN, awsRoleDurationSeconds, policy);

    // Verify the external calls.
    verify(retryPolicyFactory).getRetryPolicy();
    verify(stsOperations).assumeRole(any(AWSSecurityTokenServiceClient.class), eq(assumeRoleRequest));
    verifyNoMoreInteractionsHelper();

    // Validate the returned object.
    assertEquals(credentials, result);
}
 

@Test
public void testGetTemporarySecurityCredentialsMissingOptionalParameters()
{
    // Create an AWS parameters DTO without proxy settings.
    AwsParamsDto awsParamsDto = new AwsParamsDto();

    // Specify the duration, in seconds, of the role session.
    int awsRoleDurationSeconds = INTEGER_VALUE;

    // Create a retry policy.
    RetryPolicy retryPolicy =
        new RetryPolicy(PredefinedRetryPolicies.DEFAULT_RETRY_CONDITION, PredefinedRetryPolicies.DEFAULT_BACKOFF_STRATEGY, INTEGER_VALUE, true);

    // Create the expected assume role request.
    AssumeRoleRequest assumeRoleRequest =
        new AssumeRoleRequest().withRoleArn(AWS_ROLE_ARN).withRoleSessionName(SESSION_NAME).withDurationSeconds(awsRoleDurationSeconds);

    // Create AWS credentials for API authentication.
    Credentials credentials = new Credentials();
    credentials.setAccessKeyId(AWS_ASSUMED_ROLE_ACCESS_KEY);
    credentials.setSecretAccessKey(AWS_ASSUMED_ROLE_SECRET_KEY);
    credentials.setSessionToken(AWS_ASSUMED_ROLE_SESSION_TOKEN);

    // Create an assume role result.
    AssumeRoleResult assumeRoleResult = new AssumeRoleResult();
    assumeRoleResult.setCredentials(credentials);

    // Mock the external calls.
    when(retryPolicyFactory.getRetryPolicy()).thenReturn(retryPolicy);
    when(stsOperations.assumeRole(any(AWSSecurityTokenServiceClient.class), eq(assumeRoleRequest))).thenReturn(assumeRoleResult);

    // Call the method under test. Please note that we do not specify an IAM policy.
    Credentials result = stsDaoImpl.getTemporarySecurityCredentials(awsParamsDto, SESSION_NAME, AWS_ROLE_ARN, awsRoleDurationSeconds, null);

    // Verify the external calls.
    verify(retryPolicyFactory).getRetryPolicy();
    verify(stsOperations).assumeRole(any(AWSSecurityTokenServiceClient.class), eq(assumeRoleRequest));
    verifyNoMoreInteractionsHelper();

    // Validate the returned object.
    assertEquals(credentials, result);
}
 

/**
 * Resolve AWS credentials based on MFA/Assume role
 *
 * We will assume that if mfa_serial is defined, then role_arn and source_profile also has to be specified.
 *
 * Please note that Strongbox differ from the AWS CLI in the following:
 * AWS CLI: 'Note that configuration variables for using IAM roles can only be in the AWS CLI config file.'
 * Strongbox: '--assume-role' can be specified explicitly
 *
 * https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#using-aws-iam-roles
 */
private AWSCredentials assumeRole(ClientConfiguration clientConfiguration,
                                  ConfigProviderChain configProvider,
                                  ProfileIdentifier profile,
                                  RoleARN roleToAssume) {

    Optional<ProfileIdentifier> sourceProfile = configProvider.getSourceProfile(profile);
    if (!sourceProfile.isPresent()) {
        throw new IllegalStateException(String.format("'%s' must be specified when using '%s' for profile '%s'",
                AWSConfigPropertyKey.SOURCE_PROFILE,
                AWSConfigPropertyKey.ROLE_ARN,
                profile.name));
    }

    SessionCache sessionCache = new SessionCache(profile, roleToAssume);
    Optional<BasicSessionCredentials> cachedCredentials = sessionCache.load();

    if (cachedCredentials.isPresent()) {
        return cachedCredentials.get();
    } else {
        AWSCredentialsProvider staticCredentialsProvider = new AWSStaticCredentialsProvider(getStaticCredentials(configProvider, sourceProfile.get()));

        AWSSecurityTokenService client = AWSSecurityTokenServiceClientBuilder.standard()
                .withCredentials(staticCredentialsProvider)
                .withClientConfiguration(transformAndVerifyOrThrow(clientConfiguration))
                .withRegion(RegionResolver.getRegion())
                .build();

        String sessionId = String.format("strongbox-cli-session-%s", ZonedDateTime.now().toEpochSecond());

        AssumeRoleRequest request = new AssumeRoleRequest();
        request.withRoleArn(roleToAssume.toArn())
                .withRoleSessionName(sessionId);

        Optional<String> mfaSerial = configProvider.getMFASerial(profile);
        if (mfaSerial.isPresent()) {
            MFAToken mfaToken = mfaTokenSupplier.get();

            request.withSerialNumber(mfaSerial.get())
                    .withTokenCode(mfaToken.value);
        }

        AssumeRoleResult result = client.assumeRole(request);
        Credentials credentials = result.getCredentials();

        BasicSessionCredentials basicSessionCredentials = new BasicSessionCredentials(credentials.getAccessKeyId(), credentials.getSecretAccessKey(), credentials.getSessionToken());

        sessionCache.save(result.getAssumedRoleUser(),
                basicSessionCredentials,
                ZonedDateTime.ofInstant(credentials.getExpiration().toInstant(), ZoneId.of("UTC")));

        return basicSessionCredentials;
    }
}
 

public BasicSessionCredentials get()
{
    AWSCredentials baseCredentials = new BasicAWSCredentials(accessKeyId, secretAccessKey);

    List<Statement> statements = new ArrayList<>();
    acceptableUris.forEach(acceptableUri -> {
                Mode mode = acceptableUri.mode;
                String uri = acceptableUri.uri;
                if (uri.startsWith(URI_S3_PREFIX)) {
                    String s3BucketAndKeyStr = uri.substring(URI_S3_PREFIX.length());
                    String[] s3BucketAndKey = s3BucketAndKeyStr.split("/", 2);
                    statements.add(new Statement(Statement.Effect.Allow)
                            .withActions(S3Actions.ListObjects)
                            .withResources(new Resource("arn:aws:s3:::" + s3BucketAndKey[0])));
                    switch (mode) {
                        case READ:
                            statements.add(new Statement(Statement.Effect.Allow)
                                    .withActions(S3Actions.GetObject)
                                    .withResources(new Resource("arn:aws:s3:::" + s3BucketAndKeyStr + "*")));
                            break;
                        case WRITE:
                            statements.add(new Statement(Statement.Effect.Allow)
                                    .withActions(S3Actions.PutObject)
                                    .withResources(new Resource("arn:aws:s3:::" + s3BucketAndKeyStr + "*")));
                            break;
                    }
                }
                else if (uri.startsWith(URI_DYNAMODB_PREFIX)) {
                    String table = uri.substring(URI_DYNAMODB_PREFIX.length());
                    statements.add(new Statement(Statement.Effect.Allow)
                            .withActions(DynamoDBv2Actions.DescribeTable)
                            .withResources(new Resource(String.format("arn:aws:dynamodb:*:*:table/%s", table))));
                    switch (mode) {
                        case READ:
                            statements.add(new Statement(Statement.Effect.Allow)
                                    .withActions(DynamoDBv2Actions.Scan)
                                    .withResources(new Resource(String.format("arn:aws:dynamodb:*:*:table/%s", table))));
                            break;
                        case WRITE:
                            break;
                    }
                }
                else if (uri.startsWith(URI_EMR_PREFIX)) {
                    String cluster = uri.substring(URI_EMR_PREFIX.length());
                    // TODO: Grant minimum actions
                    statements.add(new Statement(Statement.Effect.Allow)
                                    .withActions(ElasticMapReduceActions.AllElasticMapReduceActions)
                                    .withResources(new Resource(String.format("arn:aws:elasticmapreduce:*:*:cluster/%s", cluster))));
                }
                else {
                    throw new IllegalArgumentException("Unexpected `uri`. uri=" + uri);
                }
            }
    );
    Policy policy = new Policy();
    policy.setStatements(statements);

    Credentials credentials;

    AWSSecurityTokenServiceClient stsClient = new AWSSecurityTokenServiceClient(baseCredentials);

    if (roleArn != null && !roleArn.isEmpty()) {
        // use STS to assume role
        AssumeRoleResult assumeResult = stsClient.assumeRole(new AssumeRoleRequest()
                .withRoleArn(roleArn)
                .withDurationSeconds(durationSeconds)
                .withRoleSessionName(sessionName)
                .withPolicy(policy.toJson()));

        credentials = assumeResult.getCredentials();
    }
    else {
        // Maybe we'd better add an option command later like `without_federated_token`
        GetFederationTokenRequest federationTokenRequest = new GetFederationTokenRequest()
                .withDurationSeconds(durationSeconds)
                .withName(sessionName)
                .withPolicy(policy.toJson());

        GetFederationTokenResult federationTokenResult =
                stsClient.getFederationToken(federationTokenRequest);

        credentials = federationTokenResult.getCredentials();
    }

    return new BasicSessionCredentials(
            credentials.getAccessKeyId(),
            credentials.getSecretAccessKey(),
            credentials.getSessionToken());
}
 

@Before
public void before() {
    awsEnvironment = new AWSEnvironment("Dev", "us-west-2");
    Mockito.when(gatekeeperAwsProperties.getSessionTimeout()).thenReturn(900000);
    Mockito.when(gatekeeperAwsProperties.getSessionTimeoutPad()).thenReturn(60000);

    List<Region> regions = new ArrayList<>();
    Region testRegion1 = new Region();
    Region testRegion2 = new Region();
    testRegion1.setName("us-west-2");
    testRegion2.setName("us-east-1");
    regions.add(testRegion1);
    regions.add(testRegion2);
    Account fakeAccount = new Account();
    fakeAccount.setAccountId("123");
    fakeAccount.setAlias("hello");
    fakeAccount.setRegions(regions);
    fakeAccount.setSdlc("Test");
    fakeAccount.setName("Test Account");

    AssumeRoleResult fakeRoleResult = new AssumeRoleResult();
    Credentials fakeFreshCredentials = new Credentials();   // ( ͡° ͜ʖ ͡°)
    fakeFreshCredentials.setAccessKeyId("testing");
    fakeFreshCredentials.setSecretAccessKey("s3cr3t");
    fakeFreshCredentials.setSessionToken("s35510nt0k3n");
    fakeRoleResult.setCredentials(fakeFreshCredentials);
    when(accountInformationService.getAccountByAlias("Dev")).thenReturn(fakeAccount);
    when(awsSecurityTokenServiceClient.assumeRole(any())).thenReturn(fakeRoleResult);
    when(awsSessionFactory.createEc2Session(any(), any())).thenReturn(amazonEC2Client);
    when(awsSessionFactory.createSsmSession(any(), any())).thenReturn(awsSimpleSystemsManagementClient);


}
 

/**
 * Returns a set of temporary security credentials (consisting of an access key ID, a secret access key, and a security token) that can be used to access
 * the specified AWS resource.
 *
 * @param awsSecurityTokenServiceClient the client for accessing the AWS Security Token Service
 * @param assumeRoleRequest the assume role request
 *
 * @return the response from the AssumeRole service method, as returned by AWS Security Token Service
 */
public AssumeRoleResult assumeRole(AWSSecurityTokenServiceClient awsSecurityTokenServiceClient, AssumeRoleRequest assumeRoleRequest);