Skip to content

Commit

Permalink
Fixed bulk indexing for IOCs.
Browse files Browse the repository at this point in the history 
Signed-off-by: AWSHurneyt <[email protected]>
  • Loading branch information
@AWSHurneyt
AWSHurneyt committed Jul 19, 2024
1 parent d676715 commit bf5864d
Show file tree
Hide file tree
Showing 2 changed files with 72 additions and 46 deletions.
102 changes: 58 additions & 44 deletions src/main/java/org/opensearch/securityanalytics/services/STIX2IOCFeedStore.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Locale;
Expand All @@ -57,6 +58,7 @@ public class STIX2IOCFeedStore implements FeedStore {
public static final String IOC_INDEX_PATTERN = IOC_INDEX_NAME_TEMPLATE + "-" + IOC_TIME_PLACEHOLDER;

private final Logger log = LogManager.getLogger(STIX2IOCFeedStore.class);
private final String newActiveIndex;

Instant startTime = Instant.now();

Expand All @@ -77,6 +79,8 @@ public STIX2IOCFeedStore(
this.saTifSourceConfig = saTifSourceConfig;
this.baseListener = listener;
batchSize = clusterService.getClusterSettings().get(SecurityAnalyticsSettings.BATCH_SIZE);
newActiveIndex = getNewActiveIndex(saTifSourceConfig.getId());
initSourceConfigIndexes();
}

@Override
Expand Down Expand Up @@ -109,39 +113,21 @@ public void storeIOCs(Map<IOC, UpdateAction> actionToIOCs) {
}

public void indexIocs(List<STIX2IOC> iocs) throws IOException {
String newActiveIndex = getNewActiveIndex(saTifSourceConfig.getId());
String iocIndexPattern = getAllIocIndexPatternById(saTifSourceConfig.getId());

initFeedIndex(newActiveIndex, ActionListener.wrap(
r -> {
// reset the store configs
if (saTifSourceConfig.getIocStoreConfig() instanceof DefaultIocStoreConfig) {
((DefaultIocStoreConfig) saTifSourceConfig.getIocStoreConfig()).getIocToIndexDetails().clear();
}

// recreate the store configs
saTifSourceConfig.getIocTypes().forEach(type -> {
if (saTifSourceConfig.getIocStoreConfig() instanceof DefaultIocStoreConfig) {
DefaultIocStoreConfig.IocToIndexDetails iocToIndexDetails =
new DefaultIocStoreConfig.IocToIndexDetails(new IOCType(type), iocIndexPattern, newActiveIndex);
((DefaultIocStoreConfig) saTifSourceConfig.getIocStoreConfig()).getIocToIndexDetails().add(iocToIndexDetails);
}
});
bulkIndexIocs(iocs, newActiveIndex);
}, e-> {
log.error("Failed to initialize the IOC index and save the IOCs", e);
baseListener.onFailure(e);
}
));
bulkIndexIocs(iocs, newActiveIndex);
}

private void bulkIndexIocs(List<STIX2IOC> iocs, String activeIndex) throws IOException {
if (iocs.isEmpty()) {
long duration = Duration.between(startTime, Instant.now()).toMillis();
STIX2IOCFetchService.STIX2IOCFetchResponse output = new STIX2IOCFetchService.STIX2IOCFetchResponse(Collections.emptyList(), duration);
baseListener.onResponse(output);
}

List<BulkRequest> bulkRequestList = new ArrayList<>();
BulkRequest bulkRequest = new BulkRequest();

for (STIX2IOC ioc : iocs) {
IndexRequest indexRequest = new IndexRequest(activeIndex)
.id(StringUtils.isBlank(ioc.getId())? UUID.randomUUID().toString() : ioc.getId())
.id(StringUtils.isBlank(ioc.getId()) ? UUID.randomUUID().toString() : ioc.getId())
.opType(DocWriteRequest.OpType.INDEX)
.source(ioc.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS));
bulkRequest.add(indexRequest);
Expand All @@ -151,8 +137,9 @@ private void bulkIndexIocs(List<STIX2IOC> iocs, String activeIndex) throws IOExc
bulkRequest = new BulkRequest();
}
}
bulkRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
bulkRequestList.add(bulkRequest);

if (!bulkRequest.requests().isEmpty()) bulkRequestList.add(bulkRequest);
if (!bulkRequestList.isEmpty()) bulkRequestList.get(bulkRequestList.size() - 1).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);

GroupedActionListener<BulkResponse> bulkResponseListener = new GroupedActionListener<>(ActionListener.wrap(bulkResponses -> {
int idx = 0;
Expand Down Expand Up @@ -195,22 +182,6 @@ public static String getNewActiveIndex(String sourceConfigId) {
.replace(IOC_TIME_PLACEHOLDER, Long.toString(Instant.now().toEpochMilli()));
}

public void initFeedIndex(String feedIndexName, ActionListener<CreateIndexResponse> listener) {
var indexRequest = new CreateIndexRequest(feedIndexName)
.mapping(iocIndexMapping())
.settings(Settings.builder().put("index.hidden", true).build());
client.admin().indices().create(indexRequest, ActionListener.wrap(
r -> {
log.info("Created system index {}", feedIndexName);
listener.onResponse(r);
},
e -> {
log.error("Failed to create system index {}", feedIndexName);
listener.onFailure(e);
}
));
}

public String iocIndexMapping() {
String iocMappingFile = "mappings/stix2_ioc_mapping.json";
try (InputStream is = getClass().getClassLoader().getResourceAsStream(iocMappingFile)) {
Expand All @@ -225,5 +196,48 @@ public String iocIndexMapping() {
public SATIFSourceConfig getSaTifSourceConfig() {
return saTifSourceConfig;
}

private void initSourceConfigIndexes() {
String iocIndexPattern = getAllIocIndexPatternById(saTifSourceConfig.getId());
initFeedIndex(newActiveIndex, ActionListener.wrap(
r -> {
// reset the store configs
if (saTifSourceConfig.getIocStoreConfig() instanceof DefaultIocStoreConfig) {
((DefaultIocStoreConfig) saTifSourceConfig.getIocStoreConfig()).getIocToIndexDetails().clear();
}

// recreate the store configs
saTifSourceConfig.getIocTypes().forEach(type -> {
if (saTifSourceConfig.getIocStoreConfig() instanceof DefaultIocStoreConfig) {
DefaultIocStoreConfig.IocToIndexDetails iocToIndexDetails =
new DefaultIocStoreConfig.IocToIndexDetails(new IOCType(type), iocIndexPattern, newActiveIndex);
((DefaultIocStoreConfig) saTifSourceConfig.getIocStoreConfig()).getIocToIndexDetails().add(iocToIndexDetails);
}
});

}, e-> {
log.error("Failed to initialize the IOC index and save the IOCs", e);
baseListener.onFailure(e);
}
));
}

private void initFeedIndex(String feedIndexName, ActionListener<CreateIndexResponse> listener) {
if (!clusterService.state().routingTable().hasIndex(newActiveIndex)) {
var indexRequest = new CreateIndexRequest(feedIndexName)
.mapping(iocIndexMapping())
.settings(Settings.builder().put("index.hidden", true).build());
client.admin().indices().create(indexRequest, ActionListener.wrap(
r -> {
log.info("Created system index {}", feedIndexName);
listener.onResponse(r);
},
e -> {
log.error("Failed to create system index {}", feedIndexName);
listener.onFailure(e);
}
));
}
}
}

16 changes: 14 additions & 2 deletions src/main/java/org/opensearch/securityanalytics/services/STIX2IOCFetchService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Collections;
Expand Down Expand Up @@ -118,20 +119,31 @@ public void downloadAndIndexIOCs(SATIFSourceConfig saTifSourceConfig, ActionList
STIX2IOCFeedStore feedStore = new STIX2IOCFeedStore(client, clusterService, saTifSourceConfig, listener);
STIX2IOCConsumer consumer = new STIX2IOCConsumer(batchSize, feedStore, UpdateType.REPLACE);

Instant startTime = Instant.now();
Instant endTime;
try {
log.info("Started IOC download step at {}.", startTime);
s3Connector.load(consumer);
} catch (Exception e) {
log.error("Failed to download IOCs.", e);
endTime = Instant.now();
log.error("Failed to download IOCs after {} milliseconds.", Duration.between(startTime, endTime).toMillis(), e);
listener.onFailure(e);
return;
}
endTime = Instant.now();
log.info("IOC load step took {} milliseconds.", Duration.between(startTime, endTime).toMillis());

startTime = Instant.now();
try {
log.info("Started IOC flush at {}.", startTime);
consumer.flushIOCs();
} catch (Exception e) {
log.error("Failed to flush IOCs queue.", e);
endTime = Instant.now();
log.error("Failed to flush IOCs queue after {} milliseconds.", Duration.between(startTime, endTime).toMillis(), e);
listener.onFailure(e);
}
endTime = Instant.now();
log.info("IOC flush step took {} milliseconds.", Duration.between(startTime, endTime).toMillis());
}

public void testS3Connection(S3ConnectorConfig s3ConnectorConfig, ActionListener<TestS3ConnectionResponse> listener) {
Expand Down

0 comments on commit bf5864d

Please sign in to comment.