Add helper classes for S3 connector

Signed-off-by: Chase Engelbrecht <[email protected]>

tif/build.gradle

@@ -0,0 +1,34 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+plugins {
+    id 'java'
+    id 'java-library'
+}
+
+repositories {
+    mavenLocal()
+    mavenCentral()
+    maven { url "https://aws.oss.sonatype.org/content/repositories/snapshots" }
+}
+
+dependencies {
+    implementation 'software.amazon.awssdk:s3:2.25.42'
+    implementation 'software.amazon.awssdk:sts:2.25.42'
+    implementation 'com.fasterxml.jackson.core:jackson-core:2.17.0'
+    implementation 'com.fasterxml.jackson.core:jackson-databind:2.17.0'
+    implementation 'com.fasterxml.jackson.module:jackson-module-afterburner:2.17.0'
+    implementation 'com.google.guava:guava:33.1.0-jre'
+
+    testImplementation "org.mockito:mockito-inline:5.2.0"
+    testImplementation "org.mockito:mockito-core:5.11.0"
+    testImplementation 'org.junit.jupiter:junit-jupiter-api:5.8.1'
+    testImplementation 'org.junit.jupiter:junit-jupiter-engine:5.8.1'
+    testImplementation 'org.junit.jupiter:junit-jupiter-params:5.8.1'
+}
+
+test {
+    useJUnitPlatform()
+}

tif/src/main/java/org/opensearch/securityanalytics/connector/IOCConnector.java

@@ -0,0 +1,14 @@
+package org.opensearch.securityanalytics.connector;
+
+import org.opensearch.securityanalytics.model.IOC;
+
+import java.util.List;
+
+public interface IOCConnector {
+    /**
+     * Loads a list of IOCs from a storage system
+     *
+     * @return List<IOC> a list of the retrieved IOCs
+     */
+    List<IOC> loadIOCs();
+}

tif/src/main/java/org/opensearch/securityanalytics/connector/codec/InputCodec.java

@@ -0,0 +1,16 @@
+package org.opensearch.securityanalytics.connector.codec;
+
+import org.opensearch.securityanalytics.model.IOC;
+
+import java.io.InputStream;
+import java.util.List;
+
+public interface InputCodec {
+    /**
+     * Parses an {@link InputStream} into the provided type.
+     *
+     * @param inputStream The input stream for code to process
+     * @return List<IOC> A list of IOCs parsed from the input stream
+     */
+    List<IOC> parse(InputStream inputStream);
+}

tif/src/main/java/org/opensearch/securityanalytics/connector/codec/NewlineDelimitedJsonCodec.java

@@ -0,0 +1,36 @@
+package org.opensearch.securityanalytics.connector.codec;
+
+import com.fasterxml.jackson.databind.MappingIterator;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.ObjectReader;
+import com.fasterxml.jackson.module.afterburner.AfterburnerModule;
+import com.google.common.annotations.VisibleForTesting;
+import org.opensearch.securityanalytics.exceptions.ConnectorParsingException;
+import org.opensearch.securityanalytics.model.IOC;
+
+import java.io.InputStream;
+import java.util.List;
+
+public class NewlineDelimitedJsonCodec implements InputCodec {
+    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper().registerModule(new AfterburnerModule());
+    private final ObjectReader objectReader;
+
+    public NewlineDelimitedJsonCodec(final Class<? extends IOC> clazz) {
+        this.objectReader = OBJECT_MAPPER.readerFor(clazz);
+    }
+
+    @VisibleForTesting
+    NewlineDelimitedJsonCodec(final ObjectReader objectReader) {
+        this.objectReader = objectReader;
+    }
+
+    @Override
+    public List<IOC> parse(final InputStream inputStream) {
+        try {
+            final MappingIterator<IOC> mappingIterator = objectReader.readValues(inputStream);
+            return mappingIterator.readAll();
+        } catch (final Exception e) {
+            throw new ConnectorParsingException(e);
+        }
+    }
+}

tif/src/main/java/org/opensearch/securityanalytics/connector/factory/InputCodecFactory.java

@@ -0,0 +1,13 @@
+package org.opensearch.securityanalytics.connector.factory;
+
+import org.opensearch.securityanalytics.connector.codec.InputCodec;
+import org.opensearch.securityanalytics.factory.BinaryParameterCachingFactory;
+import org.opensearch.securityanalytics.model.IOCSchema;
+import org.opensearch.securityanalytics.connector.model.InputCodecSchema;
+
+public class InputCodecFactory extends BinaryParameterCachingFactory<InputCodecSchema, IOCSchema, InputCodec> {
+    @Override
+    protected InputCodec doCreate(final InputCodecSchema inputCodecSchema, final IOCSchema iocSchema) {
+        return inputCodecSchema.getInputCodecConstructor().apply(iocSchema);
+    }
+}

tif/src/main/java/org/opensearch/securityanalytics/connector/factory/S3ClientFactory.java

@@ -0,0 +1,24 @@
+package org.opensearch.securityanalytics.connector.factory;
+
+import org.opensearch.securityanalytics.factory.BinaryParameterCachingFactory;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.s3.S3Client;
+
+public class S3ClientFactory extends BinaryParameterCachingFactory<String, String, S3Client> {
+    private final StsAssumeRoleCredentialsProviderFactory stsAssumeRoleCredentialsProviderFactory;
+
+    public S3ClientFactory(final StsAssumeRoleCredentialsProviderFactory stsAssumeRoleCredentialsProviderFactory) {
+        super();
+        this.stsAssumeRoleCredentialsProviderFactory = stsAssumeRoleCredentialsProviderFactory;
+    }
+
+    @Override
+    protected S3Client doCreate(final String roleArn, final String region) {
+        final AwsCredentialsProvider credentialsProvider = stsAssumeRoleCredentialsProviderFactory.create(roleArn, region);
+        return S3Client.builder()
+                .credentialsProvider(credentialsProvider)
+                .region(Region.of(region))
+                .build();
+    }
+}

tif/src/main/java/org/opensearch/securityanalytics/connector/factory/StsAssumeRoleCredentialsProviderFactory.java

@@ -0,0 +1,57 @@
+package org.opensearch.securityanalytics.connector.factory;
+
+import org.opensearch.securityanalytics.factory.BinaryParameterCachingFactory;
+import software.amazon.awssdk.arns.Arn;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.services.sts.StsClient;
+import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
+import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
+
+import java.util.Optional;
+import java.util.UUID;
+
+public class StsAssumeRoleCredentialsProviderFactory extends BinaryParameterCachingFactory<String, String, AwsCredentialsProvider> {
+    private static final String AWS_IAM = "iam";
+    private static final String AWS_IAM_ROLE = "role";
+
+    private final StsClientFactory stsClientFactory;
+
+    public StsAssumeRoleCredentialsProviderFactory(final StsClientFactory stsClientFactory) {
+        super();
+        this.stsClientFactory = stsClientFactory;
+    }
+
+    @Override
+    protected AwsCredentialsProvider doCreate(final String stsRoleArn, final String region) {
+        validateStsRoleArn(stsRoleArn);
+
+        final AssumeRoleRequest.Builder assumeRoleRequestBuilder = AssumeRoleRequest.builder()
+                .roleSessionName("TIF-" + UUID.randomUUID())
+                .roleArn(stsRoleArn);
+        final StsClient stsClient = stsClientFactory.create(region);
+
+        return StsAssumeRoleCredentialsProvider.builder()
+                .stsClient(stsClient)
+                .refreshRequest(assumeRoleRequestBuilder.build())
+                .build();
+    }
+
+    private void validateStsRoleArn(final String stsRoleArn) {
+        final Arn arn = getArn(stsRoleArn);
+        if (!AWS_IAM.equals(arn.service())) {
+            throw new IllegalArgumentException("roleArn must be an IAM Role");
+        }
+        final Optional<String> resourceType = arn.resource().resourceType();
+        if (resourceType.isEmpty() || !resourceType.get().equals(AWS_IAM_ROLE)) {
+            throw new IllegalArgumentException("roleArn must be an IAM Role");
+        }
+    }
+
+    private Arn getArn(final String stsRoleArn) {
+        try {
+            return Arn.fromString(stsRoleArn);
+        } catch (final Exception e) {
+            throw new IllegalArgumentException(String.format("Invalid ARN format for roleArn. Check the format of %s", stsRoleArn));
+        }
+    }
+}

tif/src/main/java/org/opensearch/securityanalytics/connector/factory/StsClientFactory.java

@@ -0,0 +1,14 @@
+package org.opensearch.securityanalytics.connector.factory;
+
+import org.opensearch.securityanalytics.factory.UnaryParameterCachingFactory;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.sts.StsClient;
+
+public class StsClientFactory extends UnaryParameterCachingFactory<String, StsClient> {
+
+    protected StsClient doCreate(final String region) {
+        return StsClient.builder()
+                .region(Region.of(region))
+                .build();
+    }
+}

tif/src/main/java/org/opensearch/securityanalytics/connector/model/InputCodecSchema.java

@@ -0,0 +1,21 @@
+package org.opensearch.securityanalytics.connector.model;
+
+import org.opensearch.securityanalytics.connector.codec.InputCodec;
+import org.opensearch.securityanalytics.connector.codec.NewlineDelimitedJsonCodec;
+import org.opensearch.securityanalytics.model.IOCSchema;
+
+import java.util.function.Function;
+
+public enum InputCodecSchema {
+    ND_JSON(iocSchema -> new NewlineDelimitedJsonCodec(iocSchema.getModelClass()));
+
+    private final Function<IOCSchema, InputCodec> inputCodecConstructor;
+
+    InputCodecSchema(final Function<IOCSchema, InputCodec> inputCodecConstructor) {
+        this.inputCodecConstructor = inputCodecConstructor;
+    }
+
+    public Function<IOCSchema, InputCodec> getInputCodecConstructor() {
+        return inputCodecConstructor;
+    }
+}

tif/src/main/java/org/opensearch/securityanalytics/exceptions/ConnectorParsingException.java

@@ -0,0 +1,7 @@
+package org.opensearch.securityanalytics.exceptions;
+
+public class ConnectorParsingException extends RuntimeException {
+    public ConnectorParsingException(final Throwable cause) {
+        super(cause);
+    }
+}

tif/src/main/java/org/opensearch/securityanalytics/factory/BinaryParameterCachingFactory.java

@@ -0,0 +1,22 @@
+package org.opensearch.securityanalytics.factory;
+
+import com.google.common.collect.HashBasedTable;
+import com.google.common.collect.Table;
+
+public abstract class BinaryParameterCachingFactory<T, U, V> {
+    private final Table<T, U, V> cache;
+
+    public BinaryParameterCachingFactory() {
+        this.cache = HashBasedTable.create();
+    }
+
+    protected abstract V doCreate(T parameter1, U parameter2);
+
+    public V create(final T parameter1, final U parameter2) {
+        if (!cache.contains(parameter1, parameter2)) {
+            cache.put(parameter1, parameter2, doCreate(parameter1, parameter2));
+        }
+
+        return cache.get(parameter1, parameter2);
+    }
+}

tif/src/main/java/org/opensearch/securityanalytics/factory/UnaryParameterCachingFactory.java

@@ -0,0 +1,22 @@
+package org.opensearch.securityanalytics.factory;
+
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+
+public abstract class UnaryParameterCachingFactory<T, U> {
+    private final Cache<T, U> cache;
+
+    public UnaryParameterCachingFactory() {
+        this.cache = CacheBuilder.newBuilder().build();
+    }
+
+    protected abstract U doCreate(T parameter);
+
+    public U create(T parameter) {
+        if (cache.getIfPresent(parameter) == null) {
+            cache.put(parameter, doCreate(parameter));
+        }
+
+        return cache.getIfPresent(parameter);
+    }
+}

tif/src/main/java/org/opensearch/securityanalytics/feed/FeedManager.java

@@ -0,0 +1,10 @@
+package org.opensearch.securityanalytics.feed;
+
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+
+public class FeedManager {
+    private static final ScheduledExecutorService EXECUTOR_SERVICE = Executors.newSingleThreadScheduledExecutor();
+
+    public 
+}

tif/src/main/java/org/opensearch/securityanalytics/feed/retriever/FeedRetriever.java

@@ -0,0 +1,18 @@
+package org.opensearch.securityanalytics.feed.retriever;
+
+import org.opensearch.securityanalytics.connector.IOCConnector;
+import org.opensearch.securityanalytics.feed.store.FeedStore;
+import org.opensearch.securityanalytics.model.IOC;
+
+import java.util.List;
+
+public class FeedRetriever implements Runnable {
+    private IOCConnector iocConnector;
+    private FeedStore feedStore;
+
+    @Override
+    public void run() {
+        final List<IOC> iocs = iocConnector.loadIOCs();
+        feedStore.storeIOCs(iocs);
+    }
+}

tif/src/main/java/org/opensearch/securityanalytics/feed/store/FeedStore.java

@@ -0,0 +1,14 @@
+package org.opensearch.securityanalytics.feed.store;
+
+import org.opensearch.securityanalytics.model.IOC;
+
+import java.util.List;
+
+public interface FeedStore {
+    /**
+     * Accepts a list of IOCs and stores them locally for use in feed processing
+     *
+     * @param iocs - A list of the IOCs to store
+     */
+    void storeIOCs(final List<IOC> iocs);
+}

tif/src/main/java/org/opensearch/securityanalytics/model/IOC.java

@@ -0,0 +1,6 @@
+package org.opensearch.securityanalytics.model;
+
+import java.io.Serializable;
+
+public abstract class IOC implements Serializable {
+}

tif/src/main/java/org/opensearch/securityanalytics/model/IOCSchema.java

@@ -0,0 +1,15 @@
+package org.opensearch.securityanalytics.model;
+
+public enum IOCSchema {
+    STIX2(STIX2.class);
+
+    private final Class<? extends IOC> modelClass;
+
+    IOCSchema(final Class<? extends IOC> modelClass) {
+        this.modelClass = modelClass;
+    }
+
+    public Class<? extends IOC> getModelClass() {
+        return modelClass;
+    }
+}

tif/src/main/java/org/opensearch/securityanalytics/model/STIX2.java

@@ -0,0 +1,36 @@
+package org.opensearch.securityanalytics.model;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+@JsonIgnoreProperties(ignoreUnknown = true)
+public class STIX2 extends IOC {
+    private String type;
+    @JsonProperty("spec_version")
+    private String specVersion;
+    private String id;
+
+    public String getType() {
+        return type;
+    }
+
+    public void setType(final String type) {
+        this.type = type;
+    }
+
+    public String getSpecVersion() {
+        return specVersion;
+    }
+
+    public void setSpecVersion(final String specVersion) {
+        this.specVersion = specVersion;
+    }
+
+    public String getId() {
+        return id;
+    }
+
+    public void setId(final String id) {
+        this.id = id;
+    }
+}