Skip to content

Commit

Permalink
make threat intel async (#703)
Browse files Browse the repository at this point in the history 
Signed-off-by: Subhobrata Dey <[email protected]>
(cherry picked from commit 0dd9787)
  • Loading branch information
@sbcd90 @github-actions
sbcd90 authored and github-actions[bot] committed Oct 31, 2023
1 parent c605fa0 commit 3f01502
Show file tree
Hide file tree
Showing 12 changed files with 534 additions and 304 deletions.
9 changes: 0 additions & 9 deletions src/main/java/org/opensearch/securityanalytics/threatIntel/DetectorThreatIntelService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,6 @@ private static String constructId(Detector detector, String iocType) {

/** Updates all detectors having threat intel detection enabled with the latest threat intel feed data*/
public void updateDetectorsWithLatestThreatIntelRules() {
try {
QueryBuilder queryBuilder =
QueryBuilders.nestedQuery("detector",
QueryBuilders.boolQuery().must(
Expand All @@ -168,7 +167,6 @@ public void updateDetectorsWithLatestThreatIntelRules() {
SearchSourceBuilder ssb = searchRequest.source();
ssb.query(queryBuilder);
ssb.size(9999);
CountDownLatch countDownLatch = new CountDownLatch(1);
client.execute(SearchDetectorAction.INSTANCE, new SearchDetectorRequest(searchRequest),
ActionListener.wrap(searchResponse -> {
List<Detector> detectors = getDetectors(searchResponse, xContentRegistry);
Expand All @@ -181,22 +179,15 @@ public void updateDetectorsWithLatestThreatIntelRules() {
ActionListener.wrap(
indexDetectorResponse -> {
log.debug("updated {} with latest threat intel info", indexDetectorResponse.getDetector().getId());
countDownLatch.countDown();
},
e -> {
log.error(() -> new ParameterizedMessage("Failed to update detector {} with latest threat intel info", detector.getId()), e);
countDownLatch.countDown();
}));
}
);
}, e -> {
log.error("Failed to fetch detectors to update with threat intel queries.", e);
countDownLatch.countDown();
}));
countDownLatch.await(5, TimeUnit.MINUTES);
} catch (InterruptedException e) {
log.error("");
}


}
Expand Down
137 changes: 99 additions & 38 deletions src/main/java/org/opensearch/securityanalytics/threatIntel/ThreatIntelFeedDataService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,15 @@
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.OpenSearchException;
import org.opensearch.OpenSearchStatusException;
import org.opensearch.action.DocWriteRequest;
import org.opensearch.action.admin.indices.create.CreateIndexRequest;
import org.opensearch.action.admin.indices.create.CreateIndexResponse;
import org.opensearch.action.bulk.BulkRequest;
import org.opensearch.action.bulk.BulkResponse;
import org.opensearch.action.index.IndexRequest;
import org.opensearch.action.search.SearchRequest;
import org.opensearch.action.support.GroupedActionListener;
import org.opensearch.action.support.IndicesOptions;
import org.opensearch.action.support.WriteRequest;
import org.opensearch.action.support.master.AcknowledgedResponse;
Expand All @@ -33,6 +36,7 @@
import org.opensearch.securityanalytics.model.ThreatIntelFeedData;
import org.opensearch.securityanalytics.threatIntel.action.PutTIFJobAction;
import org.opensearch.securityanalytics.threatIntel.action.PutTIFJobRequest;
import org.opensearch.securityanalytics.threatIntel.action.ThreatIntelIndicesResponse;
import org.opensearch.securityanalytics.threatIntel.common.TIFMetadata;
import org.opensearch.securityanalytics.threatIntel.common.StashedThreadContext;
import org.opensearch.securityanalytics.settings.SecurityAnalyticsSettings;
Expand All @@ -47,6 +51,7 @@
import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -102,17 +107,17 @@ public void getThreatIntelFeedData(

String tifdIndex = getLatestIndexByCreationDate();
if (tifdIndex == null) {
createThreatIntelFeedData();
tifdIndex = getLatestIndexByCreationDate();
createThreatIntelFeedData(listener);
} else {
SearchRequest searchRequest = new SearchRequest(tifdIndex);
searchRequest.source().size(9999); //TODO: convert to scroll
String finalTifdIndex = tifdIndex;
client.search(searchRequest, ActionListener.wrap(r -> listener.onResponse(ThreatIntelFeedDataUtils.getTifdList(r, xContentRegistry)), e -> {
log.error(String.format(
"Failed to fetch threat intel feed data from system index %s", finalTifdIndex), e);
listener.onFailure(e);
}));
}
SearchRequest searchRequest = new SearchRequest(tifdIndex);
searchRequest.source().size(9999); //TODO: convert to scroll
String finalTifdIndex = tifdIndex;
client.search(searchRequest, ActionListener.wrap(r -> listener.onResponse(ThreatIntelFeedDataUtils.getTifdList(r, xContentRegistry)), e -> {
log.error(String.format(
"Failed to fetch threat intel feed data from system index %s", finalTifdIndex), e);
listener.onFailure(e);
}));
} catch (InterruptedException e) {
log.error("Failed to get threat intel feed data", e);
listener.onFailure(e);
Expand All @@ -136,15 +141,30 @@ private String getLatestIndexByCreationDate() {
*
* @param indexName index name
*/
public void createIndexIfNotExists(final String indexName) {
public void createIndexIfNotExists(final String indexName, final ActionListener<CreateIndexResponse> listener) {
if (clusterService.state().metadata().hasIndex(indexName) == true) {
listener.onResponse(new CreateIndexResponse(true, true, indexName));
return;
}
final CreateIndexRequest createIndexRequest = new CreateIndexRequest(indexName).settings(INDEX_SETTING_TO_CREATE)
.mapping(getIndexMapping());
.mapping(getIndexMapping()).timeout(clusterSettings.get(SecurityAnalyticsSettings.THREAT_INTEL_TIMEOUT));
StashedThreadContext.run(
client,
() -> client.admin().indices().create(createIndexRequest).actionGet(clusterSettings.get(SecurityAnalyticsSettings.THREAT_INTEL_TIMEOUT))
() -> client.admin().indices().create(createIndexRequest, new ActionListener<>() {
@Override
public void onResponse(CreateIndexResponse response) {
if (response.isAcknowledged()) {
listener.onResponse(response);
} else {
onFailure(new OpenSearchStatusException("Threat intel feed index creation failed", RestStatus.INTERNAL_SERVER_ERROR));
}
}

@Override
public void onFailure(Exception e) {
listener.onFailure(e);
}
})
);
}

Expand All @@ -159,16 +179,20 @@ public void parseAndSaveThreatIntelFeedDataCSV(
final String indexName,
final Iterator<CSVRecord> iterator,
final Runnable renewLock,
final TIFMetadata tifMetadata
final TIFMetadata tifMetadata,
final ActionListener<ThreatIntelIndicesResponse> listener
) throws IOException {
if (indexName == null || iterator == null || renewLock == null) {
throw new IllegalArgumentException("Parameters cannot be null, failed to save threat intel feed data");
}

TimeValue timeout = clusterSettings.get(SecurityAnalyticsSettings.THREAT_INTEL_TIMEOUT);
Integer batchSize = clusterSettings.get(SecurityAnalyticsSettings.BATCH_SIZE);
final BulkRequest bulkRequest = new BulkRequest();

List<BulkRequest> bulkRequestList = new ArrayList<>();
BulkRequest bulkRequest = new BulkRequest();
bulkRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);

List<ThreatIntelFeedData> tifdList = new ArrayList<>();
while (iterator.hasNext()) {
CSVRecord record = iterator.next();
Expand All @@ -192,10 +216,39 @@ public void parseAndSaveThreatIntelFeedDataCSV(
bulkRequest.add(indexRequest);

if (bulkRequest.requests().size() == batchSize) {
saveTifds(bulkRequest, timeout);
bulkRequestList.add(bulkRequest);
bulkRequest = new BulkRequest();
bulkRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
}
}
bulkRequestList.add(bulkRequest);

GroupedActionListener<BulkResponse> bulkResponseListener = new GroupedActionListener<>(new ActionListener<>() {
@Override
public void onResponse(Collection<BulkResponse> bulkResponses) {
int idx = 0;
for (BulkResponse response: bulkResponses) {
BulkRequest request = bulkRequestList.get(idx);
if (response.hasFailures()) {
throw new OpenSearchException(
"error occurred while ingesting threat intel feed data in {} with an error {}",
StringUtils.join(request.getIndices()),
response.buildFailureMessage()
);
}
}
listener.onResponse(new ThreatIntelIndicesResponse(true, List.of(indexName)));
}

@Override
public void onFailure(Exception e) {
listener.onFailure(e);
}
}, bulkRequestList.size());

for (int i = 0; i < bulkRequestList.size(); ++i) {
saveTifds(bulkRequestList.get(i), timeout, bulkResponseListener);
}
saveTifds(bulkRequest, timeout);
renewLock.run();
}

Expand All @@ -206,19 +259,9 @@ public static boolean isValidIp(String ip) {
return matcher.matches();
}

public void saveTifds(BulkRequest bulkRequest, TimeValue timeout) {
public void saveTifds(BulkRequest bulkRequest, TimeValue timeout, ActionListener<BulkResponse> listener) {
try {
BulkResponse response = StashedThreadContext.run(client, () -> {
return client.bulk(bulkRequest).actionGet(timeout);
});
if (response.hasFailures()) {
throw new OpenSearchException(
"error occurred while ingesting threat intel feed data in {} with an error {}",
StringUtils.join(bulkRequest.getIndices()),
response.buildFailureMessage()
);
}
bulkRequest.requests().clear();
StashedThreadContext.run(client, () -> client.bulk(bulkRequest, listener));
} catch (OpenSearchException e) {
log.error("failed to save threat intel feed data", e);
}
Expand All @@ -241,31 +284,49 @@ public void deleteThreatIntelDataIndex(final List<String> indices) {
);
}

AcknowledgedResponse response = StashedThreadContext.run(
StashedThreadContext.run(
client,
() -> client.admin()
.indices()
.prepareDelete(indices.toArray(new String[0]))
.setIndicesOptions(IndicesOptions.LENIENT_EXPAND_OPEN_CLOSED_HIDDEN)
.execute()
.actionGet(clusterSettings.get(SecurityAnalyticsSettings.THREAT_INTEL_TIMEOUT))
);
.setTimeout(clusterSettings.get(SecurityAnalyticsSettings.THREAT_INTEL_TIMEOUT))
.execute(new ActionListener<>() {
@Override
public void onResponse(AcknowledgedResponse response) {
if (response.isAcknowledged() == false) {
onFailure(new OpenSearchException("failed to delete data[{}]", String.join(",", indices)));
}
}

if (response.isAcknowledged() == false) {
throw new OpenSearchException("failed to delete data[{}]", String.join(",", indices));
}
@Override
public void onFailure(Exception e) {
log.error("unknown exception:", e);
}
})
);
}

private void createThreatIntelFeedData() throws InterruptedException {
private void createThreatIntelFeedData(ActionListener<List<ThreatIntelFeedData>> listener) throws InterruptedException {
CountDownLatch countDownLatch = new CountDownLatch(1);
client.execute(
PutTIFJobAction.INSTANCE,
new PutTIFJobRequest("feed_updater", clusterSettings.get(SecurityAnalyticsSettings.TIF_UPDATE_INTERVAL)),
new ActionListener<AcknowledgedResponse>() {
new ActionListener<>() {
@Override
public void onResponse(AcknowledgedResponse acknowledgedResponse) {
log.debug("Acknowledged threat intel feed updater job created");
countDownLatch.countDown();
String tifdIndex = getLatestIndexByCreationDate();

SearchRequest searchRequest = new SearchRequest(tifdIndex);
searchRequest.source().size(9999); //TODO: convert to scroll
String finalTifdIndex = tifdIndex;
client.search(searchRequest, ActionListener.wrap(r -> listener.onResponse(ThreatIntelFeedDataUtils.getTifdList(r, xContentRegistry)), e -> {
log.error(String.format(
"Failed to fetch threat intel feed data from system index %s", finalTifdIndex), e);
listener.onFailure(e);
}));
}

@Override
Expand Down
43 changes: 43 additions & 0 deletions .../java/org/opensearch/securityanalytics/threatIntel/action/ThreatIntelIndicesResponse.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/
package org.opensearch.securityanalytics.threatIntel.action;

import org.opensearch.core.action.ActionResponse;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;

import java.io.IOException;
import java.util.List;

public class ThreatIntelIndicesResponse extends ActionResponse {

private Boolean isAcknowledged;

private List<String> indices;

public ThreatIntelIndicesResponse(Boolean isAcknowledged, List<String> indices) {
super();
this.isAcknowledged = isAcknowledged;
this.indices = indices;
}

public ThreatIntelIndicesResponse(StreamInput sin) throws IOException {
this(sin.readBoolean(), sin.readStringList());
}

@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeBoolean(isAcknowledged);
out.writeStringCollection(indices);
}

public Boolean isAcknowledged() {
return isAcknowledged;
}

public List<String> getIndices() {
return indices;
}
}
35 changes: 23 additions & 12 deletions ...in/java/org/opensearch/securityanalytics/threatIntel/action/TransportPutTIFJobAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.OpenSearchStatusException;
import org.opensearch.ResourceAlreadyExistsException;
import org.opensearch.action.StepListener;
import org.opensearch.action.index.IndexResponse;
Expand Down Expand Up @@ -127,12 +128,22 @@ protected ActionListener<IndexResponse> postIndexingTifJobParameter(
@Override
public void onResponse(final IndexResponse indexResponse) {
AtomicReference<LockModel> lockReference = new AtomicReference<>(lock);
try {
createThreatIntelFeedData(tifJobParameter, lockService.getRenewLockRunnable(lockReference));
} finally {
lockService.releaseLock(lockReference.get());
}
listener.onResponse(new AcknowledgedResponse(true));
createThreatIntelFeedData(tifJobParameter, lockService.getRenewLockRunnable(lockReference), new ActionListener<>() {
@Override
public void onResponse(ThreatIntelIndicesResponse threatIntelIndicesResponse) {
if (threatIntelIndicesResponse.isAcknowledged()) {
lockService.releaseLock(lockReference.get());
listener.onResponse(new AcknowledgedResponse(true));
} else {
onFailure(new OpenSearchStatusException("creation of threat intel feed data failed", RestStatus.INTERNAL_SERVER_ERROR));
}
}

@Override
public void onFailure(Exception e) {
listener.onFailure(e);
}
});
}

@Override
Expand All @@ -149,26 +160,26 @@ public void onFailure(final Exception e) {
};
}

protected void createThreatIntelFeedData(final TIFJobParameter tifJobParameter, final Runnable renewLock) {
protected void createThreatIntelFeedData(final TIFJobParameter tifJobParameter, final Runnable renewLock, final ActionListener<ThreatIntelIndicesResponse> listener) {
if (TIFJobState.CREATING.equals(tifJobParameter.getState()) == false) {
log.error("Invalid tifJobParameter state. Expecting {} but received {}", TIFJobState.CREATING, tifJobParameter.getState());
markTIFJobAsCreateFailed(tifJobParameter);
markTIFJobAsCreateFailed(tifJobParameter, listener);
return;
}

try {
tifJobUpdateService.createThreatIntelFeedData(tifJobParameter, renewLock);
tifJobUpdateService.createThreatIntelFeedData(tifJobParameter, renewLock, listener);
} catch (Exception e) {
log.error("Failed to create tifJobParameter for {}", tifJobParameter.getName(), e);
markTIFJobAsCreateFailed(tifJobParameter);
markTIFJobAsCreateFailed(tifJobParameter, listener);
}
}

private void markTIFJobAsCreateFailed(final TIFJobParameter tifJobParameter) {
private void markTIFJobAsCreateFailed(final TIFJobParameter tifJobParameter, final ActionListener<ThreatIntelIndicesResponse> listener) {
tifJobParameter.getUpdateStats().setLastFailedAt(Instant.now());
tifJobParameter.setState(TIFJobState.CREATE_FAILED);
try {
tifJobParameterService.updateJobSchedulerParameter(tifJobParameter);
tifJobParameterService.updateJobSchedulerParameter(tifJobParameter, listener);
} catch (Exception e) {
log.error("Failed to mark tifJobParameter state as CREATE_FAILED for {}", tifJobParameter.getName(), e);
}
Expand Down
Loading

0 comments on commit 3f01502

Please sign in to comment.