Skip to content

Commit

Permalink
Add per input Assume Role capability. (#49)
Browse files Browse the repository at this point in the history
  • Loading branch information
radykal-com authored and bernd committed Oct 16, 2017
1 parent 303e2ed commit 3c6ecf7
Show file tree
Hide file tree
Showing 4 changed files with 80 additions and 7 deletions.
11 changes: 11 additions & 0 deletions pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,17 @@
</exclusion>
</exclusions>
</dependency>
<dependency>
<groupId>com.amazonaws</groupId>
<artifactId>aws-java-sdk-sts</artifactId>
<version>${aws-java-sdk.version}</version>
<exclusions>
<exclusion>
<groupId>commons-logging</groupId>
<artifactId>commons-logging</artifactId>
</exclusion>
</exclusions>
</dependency>
<dependency>
<groupId>com.amazonaws</groupId>
<artifactId>amazon-kinesis-client</artifactId>
Expand Down
47 changes: 42 additions & 5 deletions src/main/java/org/graylog/aws/auth/AWSAuthProvider.java
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,10 @@
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
import com.amazonaws.auth.AWSStaticCredentialsProvider;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider;
import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClientBuilder;
import com.amazonaws.services.securitytoken.AWSSecurityTokenService;
import com.amazonaws.services.securitytoken.model.GetCallerIdentityRequest;
import org.graylog.aws.config.AWSPluginConfiguration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand All @@ -20,20 +24,53 @@ public class AWSAuthProvider implements AWSCredentialsProvider {
private AWSCredentialsProvider credentials;

public AWSAuthProvider(AWSPluginConfiguration config) {
this(config, null, null);
this(config, null, null, null,null);
}

public AWSAuthProvider(AWSPluginConfiguration config, @Nullable String accessKey, @Nullable String secretKey) {
public AWSAuthProvider(AWSPluginConfiguration config,
@Nullable String accessKey,
@Nullable String secretKey,
@Nullable String region,
@Nullable String assumeRoleArn) {
this.credentials = this.resolveAuthentication(config, accessKey, secretKey, region, assumeRoleArn);
}

private AWSCredentialsProvider resolveAuthentication(AWSPluginConfiguration config,
@Nullable String accessKey,
@Nullable String secretKey,
@Nullable String region,
@Nullable String assumeRoleArn) {
AWSCredentialsProvider awsCredentials;
if (!isNullOrEmpty(accessKey) && !isNullOrEmpty(secretKey)) {
this.credentials = new AWSStaticCredentialsProvider(new BasicAWSCredentials(accessKey, secretKey));
awsCredentials = new AWSStaticCredentialsProvider(new BasicAWSCredentials(accessKey, secretKey));
LOG.debug("Using input specific config");
} else if (!isNullOrEmpty(config.accessKey()) && !isNullOrEmpty(config.secretKey())) {
this.credentials = new AWSStaticCredentialsProvider(new BasicAWSCredentials(config.accessKey(), config.secretKey()));
awsCredentials = new AWSStaticCredentialsProvider(new BasicAWSCredentials(config.accessKey(), config.secretKey()));
LOG.debug("Using AWS Plugin config");
} else {
this.credentials = new DefaultAWSCredentialsProviderChain();
awsCredentials = new DefaultAWSCredentialsProviderChain();
LOG.debug("Using Default Provider Chain");
}
if (!isNullOrEmpty(assumeRoleArn) && !isNullOrEmpty(region)) {
LOG.debug("Creating cross account assume role credentials");
return this.getSTSCredentialsProvider(awsCredentials, region, assumeRoleArn);
} else {
return awsCredentials;
}
}

private AWSCredentialsProvider getSTSCredentialsProvider(AWSCredentialsProvider awsCredentials, String region, String assumeRoleArn) {
AWSSecurityTokenService stsClient = AWSSecurityTokenServiceClientBuilder.standard()
.withRegion(region)
.withCredentials(awsCredentials)
.build();
String roleSessionName = String.format("API_KEY_%s@ACCOUNT_%s",
awsCredentials.getCredentials().getAWSAccessKeyId(),
stsClient.getCallerIdentity(new GetCallerIdentityRequest()).getAccount());
LOG.debug("Cross account role session name: " + roleSessionName);
return new STSAssumeRoleSessionCredentialsProvider.Builder(assumeRoleArn, roleSessionName)
.withStsClient(stsClient)
.build();
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ public class CloudTrailTransport extends ThrottleableTransport {
private static final String CK_SQS_NAME = "aws_sqs_queue_name";
private static final String CK_ACCESS_KEY = "aws_access_key";
private static final String CK_SECRET_KEY = "aws_secret_key";
private static final String CK_ASSUME_ROLE_ARN = "aws_assume_role_arn";

private static final Regions DEFAULT_REGION = Regions.US_EAST_1;

Expand Down Expand Up @@ -118,7 +119,9 @@ public void doLaunch(MessageInput input) throws MisfireException {
final AWSAuthProvider authProvider = new AWSAuthProvider(
config,
input.getConfiguration().getString(CK_ACCESS_KEY),
input.getConfiguration().getString(CK_SECRET_KEY)
input.getConfiguration().getString(CK_SECRET_KEY),
input.getConfiguration().getString(CK_AWS_SQS_REGION),
input.getConfiguration().getString(CK_ASSUME_ROLE_ARN)
);

subscriber = new CloudTrailSubscriber(
Expand Down Expand Up @@ -208,6 +211,14 @@ public ConfigurationRequest getRequestedConfiguration() {
TextField.Attribute.IS_PASSWORD
));

r.addField(new TextField(
CK_ASSUME_ROLE_ARN,
"AWS assume role ARN",
"",
"The role ARN with required permissions (cross account access)",
ConfigurationField.Optional.OPTIONAL
));

return r;
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ public class KinesisTransport implements Transport {
private static final String CK_AWS_REGION = "aws_region";
private static final String CK_ACCESS_KEY = "aws_access_key";
private static final String CK_SECRET_KEY = "aws_secret_key";
private static final String CK_ASSUME_ROLE_ARN = "aws_assume_role_arn";
private static final String CK_KINESIS_STREAM_NAME = "kinesis_stream_name";

private final Configuration configuration;
Expand Down Expand Up @@ -75,7 +76,12 @@ public void launch(MessageInput input) throws MisfireException {

final AWSPluginConfiguration awsConfig = clusterConfigService.getOrDefault(AWSPluginConfiguration.class,
AWSPluginConfiguration.createDefault());
AWSAuthProvider authProvider = new AWSAuthProvider(awsConfig, configuration.getString(CK_ACCESS_KEY), configuration.getString(CK_SECRET_KEY));
AWSAuthProvider authProvider = new AWSAuthProvider(
awsConfig, configuration.getString(CK_ACCESS_KEY),
configuration.getString(CK_SECRET_KEY),
configuration.getString(CK_AWS_REGION),
configuration.getString(CK_ASSUME_ROLE_ARN)
);

this.reader = new KinesisConsumer(
configuration.getString(CK_KINESIS_STREAM_NAME),
Expand Down Expand Up @@ -159,6 +165,14 @@ public ConfigurationRequest getRequestedConfiguration() {
TextField.Attribute.IS_PASSWORD
));

r.addField(new TextField(
CK_ASSUME_ROLE_ARN,
"AWS assume role ARN",
"",
"Role ARN with required permissions (cross account access)",
ConfigurationField.Optional.OPTIONAL
));

r.addField(new TextField(
CK_KINESIS_STREAM_NAME,
"Kinesis Stream name",
Expand Down

0 comments on commit 3c6ecf7

Please sign in to comment.