CloudStore { public AWSTemporaryCredentials assumeAWSRole(String account, String roleName, String principal, Integer durationSeconds, String externalId) { if (!awsEnabled) { throw new ResourceException(ResourceException.INTERNAL_SERVER_ERROR, "AWS Support not enabled"); } final String cacheKey = getCacheKey(account, roleName, principal, durationSeconds, externalId); AWSTemporaryCredentials tempCreds = getCachedCreds(cacheKey, durationSeconds); if (tempCreds != null) { return tempCreds; } if (isFailedTempCredsRequest(cacheKey)) { LOGGER.error("CloudStore: assumeAWSRole - failed cached request for account {} with role: {}", account, roleName); return null; } AssumeRoleRequest req = getAssumeRoleRequest(account, roleName, principal, durationSeconds, externalId); try { AWSSecurityTokenService client = getTokenServiceClient(); AssumeRoleResult res = client.assumeRole(req); Credentials awsCreds = res.getCredentials(); tempCreds = new AWSTemporaryCredentials() .setAccessKeyId(awsCreds.getAccessKeyId()) .setSecretAccessKey(awsCreds.getSecretAccessKey()) .setSessionToken(awsCreds.getSessionToken()) .setExpiration(Timestamp.fromMillis(awsCreds.getExpiration().getTime())); } catch (AmazonServiceException ex) { LOGGER.error("CloudStore: assumeAWSRole - unable to assume role: {}, error: {}, status code: {}", req.getRoleArn(), ex.getMessage(), ex.getStatusCode()); if (ex.getStatusCode() == ResourceException.FORBIDDEN) { putInvalidCacheCreds(cacheKey); } return null; } catch (Exception ex) { LOGGER.error("CloudStore: assumeAWSRole - unable to assume role: {}, error: {}", req.getRoleArn(), ex.getMessage()); return null; } putCacheCreds(cacheKey, tempCreds); return tempCreds; } CloudStore(); void close(); void setHttpClient(HttpClient client); boolean isAwsEnabled(); AmazonS3 getS3Client(); AWSTemporaryCredentials assumeAWSRole(String account, String roleName, String principal,
                                                 Integer durationSeconds, String externalId); String getCloudAccount(String domainName);  }
@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(); }