NIFI-2799 AWS Assume Role Credentials With Proxy

This closes #1112.

Signed-off-by: James Wing <jvwing@gmail.com>

nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/credentials/provider/factory/CredentialPropertyDescriptors.java

@@ -155,4 +155,27 @@ public class CredentialPropertyDescriptors {
             .description("External ID for cross-account access. This is used in conjunction with role arn, " +
                 "role name, and optional session time out")
             .build();
+
+    /**
+     * Assume Role Proxy variables for configuring proxy to retrieve keys
+     */
+    public static final PropertyDescriptor ASSUME_ROLE_PROXY_HOST = new PropertyDescriptor.Builder()
+            .name("assume-role-proxy-host")
+            .displayName("Assume Role Proxy Host")
+            .expressionLanguageSupported(false)
+            .required(false)
+            .addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
+            .sensitive(false)
+            .description("Proxy host for cross-account access, if needed within your environment. This will configure a proxy to request for temporary access keys into another AWS account")
+            .build();
+
+    public static final PropertyDescriptor ASSUME_ROLE_PROXY_PORT = new PropertyDescriptor.Builder()
+            .name("assume-role-proxy-port")
+            .displayName("Assume Role Proxy Port")
+            .expressionLanguageSupported(false)
+            .required(false)
+            .addValidator(StandardValidators.POSITIVE_INTEGER_VALIDATOR)
+            .sensitive(false)
+            .description("Proxy pot for cross-account access, if needed within your environment. This will configure a proxy to request for temporary access keys into another AWS account")
+            .build();
 }

nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/credentials/provider/factory/strategies/AssumeRoleCredentialsStrategy.java

@@ -24,13 +24,18 @@ import org.apache.nifi.components.PropertyDescriptor;
 import org.apache.nifi.components.ValidationContext;
 import org.apache.nifi.components.ValidationResult;
 import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_ARN;
-import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_NAME;
-import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.MAX_SESSION_TIME;
 import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_EXTERNAL_ID;
+import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.MAX_SESSION_TIME;
+import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_NAME;
+import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_PROXY_PORT;
+import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_PROXY_HOST;
 import org.apache.nifi.processors.aws.credentials.provider.factory.CredentialsStrategy;
 
+import com.amazonaws.ClientConfiguration;
 import com.amazonaws.auth.AWSCredentialsProvider;
 import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider;
+import com.amazonaws.services.securitytoken.AWSSecurityTokenService;
+import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClient;
 
 
 /**
@@ -66,6 +71,16 @@ public class AssumeRoleCredentialsStrategy extends AbstractCredentialsStrategy {
         return false;
     }
 
+    public boolean proxyVariablesValidForAssumeRole(Map<PropertyDescriptor, String> properties){
+        final String assumeRoleProxyHost = properties.get(ASSUME_ROLE_PROXY_HOST);
+        final String assumeRoleProxyPort = properties.get(ASSUME_ROLE_PROXY_PORT);
+        if (assumeRoleProxyHost != null && !assumeRoleProxyHost.isEmpty()
+                && assumeRoleProxyPort != null && !assumeRoleProxyPort.isEmpty()) {
+            return true;
+        }
+        return false;
+    }
+
     @Override
     public Collection<ValidationResult> validate(final ValidationContext validationContext,
                                                  final CredentialsStrategy primaryStrategy) {
@@ -73,6 +88,8 @@ public class AssumeRoleCredentialsStrategy extends AbstractCredentialsStrategy {
         final boolean assumeRoleNameIsSet = validationContext.getProperty(ASSUME_ROLE_NAME).isSet();
         final Integer maxSessionTime = validationContext.getProperty(MAX_SESSION_TIME).asInteger();
         final boolean assumeRoleExternalIdIsSet = validationContext.getProperty(ASSUME_ROLE_EXTERNAL_ID).isSet();
+        final boolean assumeRoleProxyHostIsSet = validationContext.getProperty(ASSUME_ROLE_PROXY_HOST).isSet();
+        final boolean assumeRoleProxyPortIsSet = validationContext.getProperty(ASSUME_ROLE_PROXY_PORT).isSet();
 
         final Collection<ValidationResult> validationFailureResults  = new ArrayList<ValidationResult>();
 
@@ -96,6 +113,14 @@ public class AssumeRoleCredentialsStrategy extends AbstractCredentialsStrategy {
                     .build());
         }
 
+        // Both proxy host and proxy port are required if present
+        if (assumeRoleProxyHostIsSet ^ assumeRoleProxyPortIsSet){
+            validationFailureResults.add(new ValidationResult.Builder().input("Assume Role Proxy Host and Port")
+                    .valid(false)
+                    .explanation("Assume role with proxy requires both host and port for the proxy to be set")
+                    .build());
+        }
+
         return validationFailureResults;
     }
 
@@ -113,14 +138,27 @@ public class AssumeRoleCredentialsStrategy extends AbstractCredentialsStrategy {
         rawMaxSessionTime = (rawMaxSessionTime != null) ? rawMaxSessionTime : MAX_SESSION_TIME.getDefaultValue();
         final Integer maxSessionTime = Integer.parseInt(rawMaxSessionTime.trim());
         final String assumeRoleExternalId = properties.get(ASSUME_ROLE_EXTERNAL_ID);
+        STSAssumeRoleSessionCredentialsProvider.Builder builder;
+        ClientConfiguration config = new ClientConfiguration();
 
-        STSAssumeRoleSessionCredentialsProvider.Builder builder = new STSAssumeRoleSessionCredentialsProvider
+        // If proxy variables are set, then create Client Configuration with those values
+        if (proxyVariablesValidForAssumeRole(properties)) {
+            final String assumeRoleProxyHost = properties.get(ASSUME_ROLE_PROXY_HOST);
+            final Integer assumeRoleProxyPort = Integer.parseInt(properties.get(ASSUME_ROLE_PROXY_PORT));
+            config.withProxyHost(assumeRoleProxyHost);
+            config.withProxyPort(assumeRoleProxyPort);
+        }
+
+        AWSSecurityTokenService securityTokenService = new AWSSecurityTokenServiceClient(primaryCredentialsProvider, config);
+        builder = new STSAssumeRoleSessionCredentialsProvider
                 .Builder(assumeRoleArn, assumeRoleName)
-                .withLongLivedCredentialsProvider(primaryCredentialsProvider)
+                .withStsClient(securityTokenService)
                 .withRoleSessionDurationSeconds(maxSessionTime);
+
         if (assumeRoleExternalId != null && !assumeRoleExternalId.isEmpty()) {
             builder = builder.withExternalId(assumeRoleExternalId);
         }
+
         final AWSCredentialsProvider credsProvider = builder.build();
 
         return credsProvider;

nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/credentials/provider/service/AWSCredentialsProviderControllerService.java

@@ -35,15 +35,17 @@ import org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPro
 import org.apache.nifi.processors.aws.credentials.provider.factory.CredentialsProviderFactory;
 import org.apache.nifi.reporting.InitializationException;
 
-import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.USE_DEFAULT_CREDENTIALS;
+import com.amazonaws.auth.AWSCredentialsProvider;
+
 import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ACCESS_KEY;
-import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.SECRET_KEY;
+import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_EXTERNAL_ID;
+import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_PROXY_PORT;
+import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_PROXY_HOST;
 import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.CREDENTIALS_FILE;
 import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.PROFILE_NAME;
+import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.SECRET_KEY;
 import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.USE_ANONYMOUS_CREDENTIALS;
-import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_EXTERNAL_ID;
+import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.USE_DEFAULT_CREDENTIALS;
-
-import com.amazonaws.auth.AWSCredentialsProvider;
 
 /**
  * Implementation of AWSCredentialsProviderService interface
@@ -75,6 +77,8 @@ public class AWSCredentialsProviderControllerService extends AbstractControllerS
         props.add(ASSUME_ROLE_NAME);
         props.add(MAX_SESSION_TIME);
         props.add(ASSUME_ROLE_EXTERNAL_ID);
+        props.add(ASSUME_ROLE_PROXY_HOST);
+        props.add(ASSUME_ROLE_PROXY_PORT);
         properties = Collections.unmodifiableList(props);
     }
 
@@ -109,4 +113,4 @@ public class AWSCredentialsProviderControllerService extends AbstractControllerS
     public String toString() {
         return "AWSCredentialsProviderService[id=" + getIdentifier() + "]";
     }
-}
+}

nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/credentials/provider/factory/MockAWSProcessor.java

@@ -34,6 +34,8 @@ import static org.apache.nifi.processors.aws.credentials.provider.factory.Creden
 import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.PROFILE_NAME;
 import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.USE_ANONYMOUS_CREDENTIALS;
 import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_ARN;
+import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_PROXY_PORT;
+import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_PROXY_HOST;
 import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_NAME;
 import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.MAX_SESSION_TIME;
 import static org.apache.nifi.processors.aws.credentials.provider.factory.CredentialPropertyDescriptors.ASSUME_ROLE_EXTERNAL_ID;
@@ -57,7 +59,9 @@ public class MockAWSProcessor extends AbstractAWSCredentialsProviderProcessor<Am
             ASSUME_ROLE_ARN,
             ASSUME_ROLE_NAME,
             MAX_SESSION_TIME,
-            ASSUME_ROLE_EXTERNAL_ID
+            ASSUME_ROLE_EXTERNAL_ID,
+            ASSUME_ROLE_PROXY_HOST,
+            ASSUME_ROLE_PROXY_PORT
     );
 
     @Override

nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/credentials/provider/factory/TestCredentialsProviderFactory.java

@@ -198,4 +198,47 @@ public class TestCredentialsProviderFactory {
         assertEquals("credentials provider should be equal", ProfileCredentialsProvider.class,
                 credentialsProvider.getClass());
     }
+
+    @Test
+    public void testAssumeRoleCredentialsWithProxy() throws Throwable {
+        final TestRunner runner = TestRunners.newTestRunner(MockAWSProcessor.class);
+        runner.setProperty(CredentialPropertyDescriptors.CREDENTIALS_FILE, "src/test/resources/mock-aws-credentials.properties");
+        runner.setProperty(CredentialPropertyDescriptors.ASSUME_ROLE_ARN, "BogusArn");
+        runner.setProperty(CredentialPropertyDescriptors.ASSUME_ROLE_NAME, "BogusSession");
+        runner.setProperty(CredentialPropertyDescriptors.ASSUME_ROLE_PROXY_HOST, "proxy.company.com");
+        runner.setProperty(CredentialPropertyDescriptors.ASSUME_ROLE_PROXY_PORT, "8080");
+        runner.assertValid();
+
+        Map<PropertyDescriptor, String> properties = runner.getProcessContext().getProperties();
+        final CredentialsProviderFactory factory = new CredentialsProviderFactory();
+        final AWSCredentialsProvider credentialsProvider = factory.getCredentialsProvider(properties);
+        Assert.assertNotNull(credentialsProvider);
+        assertEquals("credentials provider should be equal", STSAssumeRoleSessionCredentialsProvider.class,
+                credentialsProvider.getClass());
+    }
+
+    @Test
+    public void testAssumeRoleMissingProxyHost() throws Throwable {
+        final TestRunner runner = TestRunners.newTestRunner(MockAWSProcessor.class);
+        runner.setProperty(CredentialPropertyDescriptors.CREDENTIALS_FILE, "src/test/resources/mock-aws-credentials.properties");
+        runner.setProperty(CredentialPropertyDescriptors.ASSUME_ROLE_PROXY_PORT, "8080");
+        runner.assertNotValid();
+    }
+
+    @Test
+    public void testAssumeRoleMissingProxyPort() throws Throwable {
+        final TestRunner runner = TestRunners.newTestRunner(MockAWSProcessor.class);
+        runner.setProperty(CredentialPropertyDescriptors.CREDENTIALS_FILE, "src/test/resources/mock-aws-credentials.properties");
+        runner.setProperty(CredentialPropertyDescriptors.ASSUME_ROLE_PROXY_HOST, "proxy.company.com");
+        runner.assertNotValid();
+    }
+
+    @Test
+    public void testAssumeRoleInvalidProxyPort() throws Throwable {
+        final TestRunner runner = TestRunners.newTestRunner(MockAWSProcessor.class);
+        runner.setProperty(CredentialPropertyDescriptors.CREDENTIALS_FILE, "src/test/resources/mock-aws-credentials.properties");
+        runner.setProperty(CredentialPropertyDescriptors.ASSUME_ROLE_PROXY_HOST, "proxy.company.com");
+        runner.setProperty(CredentialPropertyDescriptors.ASSUME_ROLE_PROXY_PORT, "notIntPort");
+        runner.assertNotValid();
+    }
 }