com.amazonaws.auth.policy.Policy Java Examples

/**
 * Gets the inline user policy.
 *
 * @param userName
 *            the user name
 * @param policyName
 *            the policy name
 * @param amazonIdentityManagement
 *            the amazon identity management
 * @return the inline user policy
 */
private static Policy getInlineUserPolicy(String userName, String policyName,
		AmazonIdentityManagement amazonIdentityManagement) {
	Policy policy = new Policy();
	try {
		GetUserPolicyRequest policyRequest = new GetUserPolicyRequest();
		policyRequest.setUserName(userName);
		policyRequest.setPolicyName(policyName);
		GetUserPolicyResult policyResult = amazonIdentityManagement.getUserPolicy(policyRequest);
		String policyAsString = policyResult.getPolicyDocument();

		policyAsString = java.net.URLDecoder.decode(policyAsString, "UTF-8");
		policy = Policy.fromJson(policyAsString);
	} catch (Exception e) {
		logger.error(e.getMessage());
	}

	return policy;
}
 

@Test
public void shouldThrowException_onAmazonClientExceptionFromSetPolicy() {
    // Given
    final Policy mockPolicy = mock(Policy.class);
    final String mockPolicyJson = randomString();
    when(mockPolicy.toJson()).thenReturn(mockPolicyJson);
    doThrow(AmazonClientException.class).when(mockAmazonSnsClient)
            .setTopicAttributes(any(SetTopicAttributesRequest.class));

    // When
    AmazonClientException thrownException = null;
    try {
        snsTopicResource.setPolicy(mockPolicy);
    } catch (final AmazonClientException e) {
        thrownException = e;
    }

    // Then
    assertNotNull(thrownException);
}
 

public boolean isTrusted(InstanceProfile instanceProfile, ValidationResultBuilder resultBuilder) {
    List<Role> instanceProfileRoles = instanceProfile.getRoles();
    for (Role role : instanceProfileRoles) {
        Policy assumeRolePolicy = awsIamService.getAssumeRolePolicy(role);
        if (assumeRolePolicy != null) {
            for (Statement statement : assumeRolePolicy.getStatements()) {
                if (checkAssumeRoleInActions(statement.getActions()) &&
                        checkEC2InPrincipals(statement.getPrincipals())) {
                    return true;
                }
            }
        }
    }
    resultBuilder.error(
            String.format("The instance profile (%s) doesn't have an EC2 trust relationship.",
                    instanceProfile.getArn()));
    return false;
}
 

/**
 * Gets the role assume role policy document as a Policy object
 *
 * @param role Role to evaluate
 * @return assume role Policy object
 */
public Policy getAssumeRolePolicy(Role role) {
    Policy policy = null;
    String assumeRolePolicyDocument = role.getAssumeRolePolicyDocument();
    if (assumeRolePolicyDocument != null) {
        try {
            String decodedAssumeRolePolicyDocument = URLDecoder.decode(assumeRolePolicyDocument,
                    StandardCharsets.UTF_8);
            policy = Policy.fromJson(decodedAssumeRolePolicyDocument);
        } catch (IllegalArgumentException e) {
            LOGGER.error(String.format("Unable to get policy from role (%s)", role.getArn()), e);
        }
    }

    return policy;
}
 

/**
 * 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 test_that_removePolicyFromStatement_removes_the_given_statement() {

  String removeId = "remove id";
  String keepId = "keep id";
  Statement statementToRemove =
      new Statement(Statement.Effect.Allow)
          .withId(removeId)
          .withActions(KMSActions.AllKMSActions);
  Statement statementToKeep =
      new Statement(Statement.Effect.Deny).withId(keepId).withActions(KMSActions.AllKMSActions);
  Policy policy = new Policy("policy", Lists.newArrayList(statementToKeep, statementToRemove));

  kmsPolicyService.removeStatementFromPolicy(policy, removeId);

  assertTrue(policy.getStatements().contains(statementToKeep));
  assertFalse(policy.getStatements().contains(statementToRemove));
}
 

/**
 * Attempts to download the policy, if something goes wrong, it returns an empty optional.
 *
 * @param kmsCmkId The KMS CMK that you want the default policy for
 * @param regionName The region that the KMS CMK resides in.
 * @return The policy if it can successfully be fetched.
 */
protected Optional<Policy> downloadPolicy(String kmsCmkId, String regionName, int retryCount) {
  final Set<String> unretryableErrors =
      ImmutableSet.of("NotFoundException", "InvalidArnException", "AccessDeniedException");
  try {
    var policy =
        kmsPolicyService.getPolicyFromPolicyString(getKmsKeyPolicy(kmsCmkId, regionName));
    return Optional.of(policy);
  } catch (AWSKMSException e) {
    String errorCode = e.getErrorCode();
    logger.error("Failed to download policy, error code: {}", errorCode);
    if (!unretryableErrors.contains(errorCode) && retryCount < 10) {
      try {
        Thread.sleep(500 ^ (retryCount + 1));
      } catch (InterruptedException e1) {
        return Optional.empty();
      }
      return downloadPolicy(kmsCmkId, regionName, retryCount + 1);
    }
  }
  return Optional.empty();
}
 

@Test
public void shouldSetPolicy_withPolicy() {
    // Given
    final Policy mockPolicy = mock(Policy.class);
    final String mockPolicyJson = randomString();
    when(mockPolicy.toJson()).thenReturn(mockPolicyJson);

    // When
    snsTopicResource.setPolicy(mockPolicy);

    // Then
    final ArgumentCaptor<SetTopicAttributesRequest> captor = ArgumentCaptor
            .forClass(SetTopicAttributesRequest.class);
    verify(mockAmazonSnsClient).setTopicAttributes(captor.capture());
    final SetTopicAttributesRequest setTopicAttributesRequest = captor.getValue();
    assertEquals(topicArn, setTopicAttributesRequest.getTopicArn());
    assertEquals("Policy", setTopicAttributesRequest.getAttributeName());
    assertEquals(mockPolicyJson, setTopicAttributesRequest.getAttributeValue());
}
 

/**
 * Gets the inline role policy.
 *
 * @param roleName
 *            the role name
 * @param policyName
 *            the policy name
 * @param amazonIdentityManagement
 *            the amazon identity management
 * @return the inline role policy
 */
private static Policy getInlineRolePolicy(String roleName, String policyName,
		AmazonIdentityManagement amazonIdentityManagement) {
	Policy policy = new Policy();
	try {
		GetRolePolicyRequest policyRequest = new GetRolePolicyRequest();
		policyRequest.setRoleName(roleName);
		policyRequest.setPolicyName(policyName);
		GetRolePolicyResult policyResult = amazonIdentityManagement.getRolePolicy(policyRequest);
		String policyAsString = policyResult.getPolicyDocument();

		policyAsString = java.net.URLDecoder.decode(policyAsString, "UTF-8");
		policy = Policy.fromJson(policyAsString);
	} catch (Exception e) {
		logger.error(e.getMessage());
	}

	return policy;
}
 

/**
 * Gets the inline role policy.
 *
 * @param roleName
 *            the role name
 * @param amazonIdentityManagement
 *            the amazon identity management
 * @param actionSet
 *            the action set
 * @return the inline role policy
 */
private static Set<String> getInlineRolePolicyActionSet(String roleName,
		AmazonIdentityManagementClient amazonIdentityManagement) {
	Set<String> actionSet = new HashSet<>();

	List<String> inlineRolePolicyNameList = new ArrayList<>();
	ListRolePoliciesRequest listRolePoliciesRequest = new ListRolePoliciesRequest();
	listRolePoliciesRequest.setRoleName(roleName);
	ListRolePoliciesResult listRolePoliciesResult = null;
	do {
		listRolePoliciesResult = amazonIdentityManagement.listRolePolicies(listRolePoliciesRequest);
		inlineRolePolicyNameList.addAll(listRolePoliciesResult.getPolicyNames());
		listRolePoliciesRequest.setMarker(listRolePoliciesResult.getMarker());
	} while (listRolePoliciesResult.isTruncated());

	for (String policyName : inlineRolePolicyNameList) {
		Policy policy = getInlineRolePolicy(roleName, policyName, amazonIdentityManagement);
		actionSet.addAll(getActionSet(policy));
	}
	return actionSet;
}
 

/**
 * Gets the inline user policy.
 *
 * @param userName
 *            the user name
 * @param amazonIdentityManagement
 *            the amazon identity management
 * @param actionSet
 *            the action set
 * @return the inline user policy
 */
private static Set<String> getInlineUserPolicyActionSet(String userName,
		AmazonIdentityManagementClient amazonIdentityManagement) {
	Set<String> actionSet = new HashSet<>();

	List<String> inlineUserPolicyNameList = new ArrayList<>();
	ListUserPoliciesRequest listUserPoliciesRequest = new ListUserPoliciesRequest();
	listUserPoliciesRequest.setUserName(userName);
	ListUserPoliciesResult listUserPoliciesResult = null;
	do {
		listUserPoliciesResult = amazonIdentityManagement.listUserPolicies(listUserPoliciesRequest);
		inlineUserPolicyNameList.addAll(listUserPoliciesResult.getPolicyNames());
		listUserPoliciesRequest.setMarker(listUserPoliciesResult.getMarker());
	} while (listUserPoliciesResult.isTruncated());

	for (String policyName : inlineUserPolicyNameList) {
		Policy policy = getInlineUserPolicy(userName, policyName, amazonIdentityManagement);
		actionSet.addAll(getActionSet(policy));
	}
	return actionSet;
}
 

@Test
public void testGetAssumeRolePolicyDocument() throws IOException {
    String assumeRolePolicyDocument = awsIamService.getResourceFileAsString(
            "json/aws-assume-role-policy-document.json");
    String encodedAssumeRolePolicyDocument = URLEncoder.encode(assumeRolePolicyDocument,
            StandardCharsets.UTF_8);


    Statement statement = new Statement(Effect.Allow).withId("1")
            .withPrincipals(new Principal("AWS", "arn:aws:iam::123456890:role/assume-role"))
            .withActions(SecurityTokenServiceActions.AssumeRole);
    Policy expectedAssumeRolePolicy = new Policy().withStatements(statement);

    Role role = mock(Role.class);
    when(role.getAssumeRolePolicyDocument()).thenReturn(encodedAssumeRolePolicyDocument);

    Policy assumeRolePolicy = awsIamService.getAssumeRolePolicy(role);
    assertThat(assumeRolePolicy).isNotNull();
    assertThat(assumeRolePolicy.toJson()).isEqualTo(expectedAssumeRolePolicy.toJson());
}
 

@Test
public void shouldSetQueueAttributes_withPolicy() {
    // Given
    final Policy mockPolicy = mock(Policy.class);
    final String mockPolicyJson = randomString();
    when(mockPolicy.toJson()).thenReturn(mockPolicyJson);

    // When
    sqsQueueResource.setPolicy(mockPolicy);

    // Then
    final ArgumentCaptor<SetQueueAttributesRequest> captor = ArgumentCaptor
            .forClass(SetQueueAttributesRequest.class);
    verify(amazonSqsClient).setQueueAttributes(captor.capture());
    final SetQueueAttributesRequest setQueueAttributesRequest = captor.getValue();
    assertEquals(queueUrl, setQueueAttributesRequest.getQueueUrl());
    assertEquals(mockPolicyJson, setQueueAttributesRequest.getAttributes()
            .get(QueueAttributeName.Policy.toString()));
}
 

@Test
public void invalidInstanceProfileTrustOneRoleNoTrustPolicy() {
    Role role = new Role().withArn("roleArn").withAssumeRolePolicyDocument(new Policy().toJson());
    InstanceProfile instanceProfile = new InstanceProfile().withArn("oneRoleNoTrustPolicy")
            .withRoles(role);
    checkInvalidInstanceProfileTrust(instanceProfile);
}
 

/**
 * For retrieving vault inventory. For initializing SQS for determining when
 * job completed. Does nothing if member snsTopicName is null. Sets members
 * sqsQueueURL, sqsQueueARN, and sqsClient.
 */
   private void setupSQS() {
	// If no sqsQueueName setup then simply return
	if (sqsQueueName == null)
		return;

	CreateQueueRequest request = new CreateQueueRequest()
			.withQueueName(sqsQueueName);
	CreateQueueResult result = sqsClient.createQueue(request);
	sqsQueueURL = result.getQueueUrl();

	GetQueueAttributesRequest qRequest = new GetQueueAttributesRequest()
			.withQueueUrl(sqsQueueURL).withAttributeNames("QueueArn");

	GetQueueAttributesResult qResult = sqsClient
			.getQueueAttributes(qRequest);
	sqsQueueARN = qResult.getAttributes().get("QueueArn");

	Policy sqsPolicy = new Policy().withStatements(new Statement(
			Effect.Allow).withPrincipals(Principal.AllUsers)
			.withActions(SQSActions.SendMessage)
			.withResources(new Resource(sqsQueueARN)));
	Map<String, String> queueAttributes = new HashMap<String, String>();
	queueAttributes.put("Policy", sqsPolicy.toJson());
	sqsClient.setQueueAttributes(new SetQueueAttributesRequest(sqsQueueURL,
			queueAttributes));
}
 

private TemporaryQueue createQueue(String snsTopicArn, String sqsQueueArn, String sqsQueueName) {
  String sqsQueueUrl =
      amazonSQS
          .createQueue(
              new CreateQueueRequest()
                  .withQueueName(sqsQueueName)
                  .withAttributes(
                      Collections.singletonMap(
                          "MessageRetentionPeriod", "60")) // 60s message retention
              )
          .getQueueUrl();
  log.info("Created Temporary S3 Notification Queue: {}", value("queue", sqsQueueUrl));

  String snsTopicSubscriptionArn =
      amazonSNS.subscribe(snsTopicArn, "sqs", sqsQueueArn).getSubscriptionArn();

  Statement snsStatement =
      new Statement(Statement.Effect.Allow).withActions(SQSActions.SendMessage);
  snsStatement.setPrincipals(Principal.All);
  snsStatement.setResources(Collections.singletonList(new Resource(sqsQueueArn)));
  snsStatement.setConditions(
      Collections.singletonList(
          new Condition()
              .withType("ArnEquals")
              .withConditionKey("aws:SourceArn")
              .withValues(snsTopicArn)));

  Policy allowSnsPolicy = new Policy("allow-sns", Collections.singletonList(snsStatement));

  HashMap<String, String> attributes = new HashMap<>();
  attributes.put("Policy", allowSnsPolicy.toJson());
  amazonSQS.setQueueAttributes(sqsQueueUrl, attributes);

  return new TemporaryQueue(snsTopicArn, sqsQueueArn, sqsQueueUrl, snsTopicSubscriptionArn);
}
 

/**
 * Validates the cloudFileSystem
 *
 * @param cloudFileSystem         cloud file system to evaluate
 * @param validationResultBuilder builder for any errors encountered
 */
public void validate(AmazonIdentityManagement iam, CloudS3View cloudFileSystem,
        ValidationResultBuilder validationResultBuilder) {
    AccountMappingBase accountMappings = cloudFileSystem.getAccountMapping();
    if (accountMappings != null) {
        SortedSet<String> roleArns = getRoleArnsForUsers(getUsers(), accountMappings.getUserMappings());
        LOGGER.info("Getting role from AWS, roleArns.size: {}, roleArns: {}", roleArns.size(), roleArns);
        Set<Role> roles = awsIamService.getValidRoles(iam, roleArns, validationResultBuilder);

        boolean s3guardEnabled = cloudFileSystem.getS3GuardDynamoTableName() != null;
        List<String> policyFileNames = getPolicyFileNames(s3guardEnabled);

        SortedSet<String> failedActions = new TreeSet<>();
        for (StorageLocationBase location : cloudFileSystem.getLocations()) {
            if (checkLocation(location)) {
                Map<String, String> replacements = getPolicyJsonReplacements(location, cloudFileSystem);
                List<Policy> policies = getPolicies(policyFileNames, replacements);
                for (Role role : roles) {
                    try {
                        List<EvaluationResult> evaluationResults = awsIamService.validateRolePolicies(iam,
                                role, policies);
                        failedActions.addAll(getFailedActions(role, evaluationResults));
                    } catch (AmazonIdentityManagementException e) {
                        // Only log the error and keep processing. Failed actions won't be added, but
                        // processing doesn't get stopped either. This can happen due to rate limiting.
                        LOGGER.error("Unable to validate role policies for role {} due to {}", role.getArn(),
                                e.getMessage(), e);
                    }
                }
            }
        }
        if (!failedActions.isEmpty()) {
            String errorMessage = String.format("The role(s) (%s) don't have the required permissions:%n%s",
                            String.join(", ", roles.stream().map(Role::getArn).collect(Collectors.toCollection(TreeSet::new))),
                            String.join("\n", failedActions));
            LOGGER.warn(errorMessage);
            validationResultBuilder.error(errorMessage);
        }
    }
}
 

/**
 * Returns Policy objects after the policy json has had replacements
 *
 * @param policyFileNames policy file names to read before replacement
 * @param replacements    simple replacements to be made to the policy json
 * @return list of Policy objects after replacements made to policy json
 */
List<Policy> getPolicies(List<String> policyFileNames,
        Map<String, String> replacements) {
    List<Policy> policies = new ArrayList<>(policyFileNames.size());
    for (String policyFileName : policyFileNames) {
        Policy policy = awsIamService.getPolicy(policyFileName, replacements);
        if (policy != null) {
            policies.add(policy);
        }
    }
    return policies;
}
 

public void validate(AmazonIdentityManagement iam, InstanceProfile instanceProfile,
        CloudS3View cloudFileSystem, ValidationResultBuilder validationResultBuilder) {
    SortedSet<String> failedActions = new TreeSet<>();

    // TODO need to figure out how to get LOGS_LOCATION_BASE value
    Map<String, String> replacements = Map.ofEntries(
            Map.entry("${LOGS_LOCATION_BASE}", "")
    );

    Policy policy = awsIamService.getPolicy("aws-cdp-log-policy.json", replacements);
    List<Role> roles = instanceProfile.getRoles();
    List<Policy> policies = Collections.singletonList(policy);
    for (Role role : roles) {
        try {
            List<EvaluationResult> evaluationResults = awsIamService.validateRolePolicies(iam,
                    role, policies);
            failedActions.addAll(getFailedActions(role, evaluationResults));
        } catch (AmazonIdentityManagementException e) {
            // Only log the error and keep processing. Failed actions won't be added, but
            // processing doesn't get stopped either. This can happen due to rate limiting.
            LOGGER.error("Unable to validate role policies for role {} due to {}", role.getArn(),
                    e.getMessage(), e);
        }
    }

    if (!failedActions.isEmpty()) {
        validationResultBuilder.error(String.format("The log role (%s) don't have the required permissions: %n%s",
                String.join(", ", roles.stream().map(Role::getArn).collect(Collectors.toCollection(TreeSet::new))),
                String.join("\n", failedActions)));
    }
}
 

/**
 * Returns a Policy object that has replacements made to the template json
 *
 * @param policyFileName Policy template file name
 * @param replacements   map of simple replacements to make to policy template json
 * @return Policy with replacements made
 */
public Policy getPolicy(String policyFileName, Map<String, String> replacements) {
    Policy policy = null;
    try {
        String policyJsonTemplate = getResourceFileAsString(POLICY_BASE_LOCATION + policyFileName);
        String policyJson = handleTemplateReplacements(policyJsonTemplate, replacements);
        if (policyJson != null) {
            policy = Policy.fromJson(policyJson);
        }
    } catch (IOException e) {
        LOGGER.debug("Unable to load policy json", e);
    }

    return policy;
}
 

/**
 * Validates the given roles against the policies
 *
 * @param iam      AmazonIdentityManagement client
 * @param role     Role object to check
 * @param policies collection of Policy objects to check
 * @return list of evaluation results
 */
public List<EvaluationResult> validateRolePolicies(AmazonIdentityManagement iam, Role role,
        Collection<Policy> policies) throws AmazonIdentityManagementException {
    List<EvaluationResult> evaluationResults = new ArrayList<>();
    for (Policy policy : policies) {
        for (Statement statement : policy.getStatements()) {
            SortedSet<String> actions = getStatementActions(statement);
            SortedSet<String> resources = getStatementResources(statement);
            List<EvaluationResult> results = simulatePrincipalPolicy(iam, role.getArn(), actions, resources);
            evaluationResults.addAll(results);
        }
    }
    return evaluationResults;
}
 

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());
}
 

@Test
public void invalidInstanceProfileTrustMultipleRolesNoTrustPolicy() {
    Role role1 = new Role().withArn("roleArn").withAssumeRolePolicyDocument(new Policy().toJson());
    Role role2 = new Role().withArn("roleArn").withAssumeRolePolicyDocument(new Policy().toJson());
    InstanceProfile instanceProfile = new InstanceProfile().withArn("multipleRolesNoTrustPolicy")
            .withRoles(role1, role2);
    checkInvalidInstanceProfileTrust(instanceProfile);
}
 

@Test
public void validInstanceProfileTrustOneRoleTrusted() {
    Policy trustedPolicy = getTrustedPolicy();
    Role role = new Role().withArn("roleArn").withAssumeRolePolicyDocument(trustedPolicy.toJson());
    InstanceProfile instanceProfile = new InstanceProfile().withArn("oneRoleTrusted")
            .withRoles(role);
    checkValidInstanceProfileTrust(instanceProfile);
}
 

@Test
public void validInstanceProfileTrustMultipleRolesTrusted() {
    Policy untrustedPolicy = new Policy();
    Role role1 = new Role().withArn("roleArn").withAssumeRolePolicyDocument(untrustedPolicy.toJson());
    Policy trustedPolicy = getTrustedPolicy();
    Role role2 = new Role().withArn("roleArn").withAssumeRolePolicyDocument(trustedPolicy.toJson());
    InstanceProfile instanceProfile = new InstanceProfile().withArn("multipleRolesTrusted")
            .withRoles(role1, role2);
    checkValidInstanceProfileTrust(instanceProfile);
}
 

private Policy getTrustedPolicy() {
    return new Policy().withStatements(
            new Statement(Effect.Allow)
                    .withActions(SecurityTokenServiceActions.AssumeRole)
                    .withPrincipals(new Principal("Service", Services.AmazonEC2.getServiceId()))
    );
}
 

@Test
public void testGetPolicy() {
    assertThat(awsIamService.getPolicy("abc", Collections.emptyMap())).isNull();

    Policy expectedPolicyNoReplacements = new Policy().withStatements(
            new Statement(Effect.Allow).withId("FullObjectAccessUnderAuditDir")
                    .withActions(S3Actions.GetObject, S3Actions.PutObject)
                    .withResources(new Resource("arn:aws:s3:::${STORAGE_LOCATION_BASE}/ranger/audit/*")),
            new Statement(Effect.Allow).withId("LimitedAccessToDataLakeBucket")
                    .withActions(S3Actions.AbortMultipartUpload, S3Actions.ListObjects,
                            S3Actions.ListBucketMultipartUploads)
                    .withResources(new Resource("arn:aws:s3:::${DATALAKE_BUCKET}"))
    );
    assertThat(awsIamService.getPolicy("aws-cdp-ranger-audit-s3-policy.json",
            Collections.emptyMap()).toJson()).isEqualTo(expectedPolicyNoReplacements.toJson());

    Policy expectedPolicyWithReplacements = new Policy().withStatements(
            new Statement(Effect.Allow).withId("FullObjectAccessUnderAuditDir")
                    .withActions(S3Actions.GetObject, S3Actions.PutObject)
                    .withResources(new Resource("arn:aws:s3:::mybucket/mycluster/ranger/audit/*")),
            new Statement(Effect.Allow).withId("LimitedAccessToDataLakeBucket")
                    .withActions(S3Actions.AbortMultipartUpload, S3Actions.ListObjects,
                            S3Actions.ListBucketMultipartUploads)
                    .withResources(new Resource("arn:aws:s3:::mybucket"))
    );

    Map<String, String> policyReplacements = new HashMap<>();
    policyReplacements.put("${STORAGE_LOCATION_BASE}", "mybucket/mycluster");
    policyReplacements.put("${DATALAKE_BUCKET}", "mybucket");
    assertThat(awsIamService.getPolicy("aws-cdp-ranger-audit-s3-policy.json",
            policyReplacements).toJson()).isEqualTo(expectedPolicyWithReplacements.toJson());
}
 

public static Policy getPolicy() {
  	Policy policy = new Policy();
List<Statement> statements = new ArrayList<Statement>();
Statement statement = new Statement(Effect.Allow);
  	List<Action> actions = new ArrayList<>();
actions.add(IdentityManagementActions.AllIdentityManagementActions);
actions.add(EC2Actions.RunInstances);
statement.setActions(actions);
statements.add(statement);
policy.setStatements(statements);
      policy.setId("123");
      policy.setStatements(statements);
return policy;
  }
 

protected Policy allowSendMessagePolicy(String roleARN) {
    Policy policy = new Policy();
    Statement statement = new Statement(Statement.Effect.Allow);
    statement.setActions(Collections.singletonList(SQSActions.SendMessage));
    statement.setPrincipals(new Principal(roleARN));
    statement.setResources(Collections.singletonList(new Resource("arn:aws:sqs:*:*:*")));
    policy.setStatements(Collections.singletonList(statement));
    return policy;
}
 

@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);
}