add field based rules support in correlation engine

src/main/java/org/opensearch/securityanalytics/correlation/JoinEngine.java

@@ -4,7 +4,8 @@
  */
 package org.opensearch.securityanalytics.correlation;
 
-import kotlin.Pair;
+import org.apache.commons.lang3.tuple.Pair;
+import org.apache.commons.lang3.tuple.Triple;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.search.join.ScoreMode;
@@ -15,6 +16,7 @@
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.client.Client;
+import org.opensearch.cluster.routing.Preference;
 import org.opensearch.common.xcontent.LoggingDeprecationHandler;
 import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.commons.alerting.action.PublishFindingsRequest;
@@ -37,11 +39,13 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 
@@ -127,25 +131,34 @@
     private void getValidDocuments(String detectorType, List<String> indices, List<CorrelationRule> correlationRules, List<String> relatedDocIds) {
         MultiSearchRequest mSearchRequest = new MultiSearchRequest();
         List<CorrelationRule> validCorrelationRules = new ArrayList<>();
+        List<String> validFields = new ArrayList<>();
 
         for (CorrelationRule rule: correlationRules) {
             Optional<CorrelationQuery> query = rule.getCorrelationQueries().stream()
                     .filter(correlationQuery -> correlationQuery.getCategory().equals(detectorType)).findFirst();
 
             if (query.isPresent()) {
                 BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery()
-                        .filter(QueryBuilders.termsQuery("_id", relatedDocIds))
-                        .must(QueryBuilders.queryStringQuery(query.get().getQuery()));
+                        .filter(QueryBuilders.termsQuery("_id", relatedDocIds));
 
+                if (query.get().getField() != null) {
+                    queryBuilder = queryBuilder.must(QueryBuilders.existsQuery(query.get().getField()));
+                } else {
+                    queryBuilder = queryBuilder.must(QueryBuilders.queryStringQuery(query.get().getQuery()));
+                }
                 SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
                 searchSourceBuilder.query(queryBuilder);
                 searchSourceBuilder.fetchSource(false);
+                if (query.get().getField() != null) {
+                    searchSourceBuilder.fetchField(query.get().getField());
+                }
                 searchSourceBuilder.size(10000);
                 SearchRequest searchRequest = new SearchRequest();
                 searchRequest.indices(indices.toArray(new String[]{}));
                 searchRequest.source(searchSourceBuilder);
 
                 validCorrelationRules.add(rule);
+                validFields.add(query.get().getField());
                 mSearchRequest.add(searchRequest);
             }
         }
@@ -155,7 +168,7 @@
                 @Override
                 public void onResponse(MultiSearchResponse items) {
                     MultiSearchResponse.Item[] responses = items.getResponses();
-                    List<CorrelationRule> filteredCorrelationRules = new ArrayList<>();
+                    List<Triple<CorrelationRule, SearchHit[], String>> filteredCorrelationRules = new ArrayList<>();
 
                     int idx = 0;
                     for (MultiSearchResponse.Item response : responses) {
@@ -165,14 +178,17 @@
                         }
 
                         if (response.getResponse().getHits().getTotalHits().value > 0L) {
-                            filteredCorrelationRules.add(validCorrelationRules.get(idx));
+                            filteredCorrelationRules.add(Triple.of(validCorrelationRules.get(idx),
+                                    response.getResponse().getHits().getHits(), validFields.get(idx)));
                         }
                         ++idx;
                     }
 
                     Map<String, List<CorrelationQuery>> categoryToQueriesMap = new HashMap<>();
-                    for (CorrelationRule rule: filteredCorrelationRules) {
-                        List<CorrelationQuery> queries = rule.getCorrelationQueries();
+                    Map<String, Long> categoryToTimeWindowMap = new HashMap<>();
+                    for (Triple<CorrelationRule, SearchHit[], String> rule: filteredCorrelationRules) {
+                        List<CorrelationQuery> queries = rule.getLeft().getCorrelationQueries();
+                        Long timeWindow = rule.getLeft().getCorrTimeWindow();
 
                         for (CorrelationQuery query: queries) {
                             List<CorrelationQuery> correlationQueries;
@@ -181,12 +197,36 @@
                             } else {
                                 correlationQueries = new ArrayList<>();
                             }
-                            correlationQueries.add(query);
+                            if (categoryToTimeWindowMap.containsKey(query.getCategory())) {
+                                categoryToTimeWindowMap.put(query.getCategory(), Math.max(timeWindow, categoryToTimeWindowMap.get(query.getCategory())));
+                            } else {
+                                categoryToTimeWindowMap.put(query.getCategory(), timeWindow);
+                            }
+
+                            if (query.getField() == null) {
+                                correlationQueries.add(query);
+                            } else {
+                                SearchHit[] hits = rule.getMiddle();
+                                StringBuilder qb = new StringBuilder(query.getField()).append(":(");
+                                for (int i = 0; i < hits.length; ++i) {
+                                    String value = hits[i].field(rule.getRight()).getValue();
+                                    qb.append(value);
+                                    if (i < hits.length-1) {
+                                        qb.append(" OR ");
+                                    } else {
+                                        qb.append(")");
+                                    }
+                                }
+                                if (query.getQuery() != null) {
+                                    qb.append(" AND ").append(query.getQuery());
+                                }
+                                correlationQueries.add(new CorrelationQuery(query.getIndex(), qb.toString(), query.getCategory(), null));
+                            }
                             categoryToQueriesMap.put(query.getCategory(), correlationQueries);
                         }
                     }
-                    searchFindingsByTimestamp(detectorType, categoryToQueriesMap,
-                            filteredCorrelationRules.stream().map(CorrelationRule::getId).collect(Collectors.toList()));
+                    searchFindingsByTimestamp(detectorType, categoryToQueriesMap, categoryToTimeWindowMap,
+                            filteredCorrelationRules.stream().map(Triple::getLeft).map(CorrelationRule::getId).collect(Collectors.toList()));
                 }
 
                 @Override
@@ -203,15 +243,15 @@
      * this method searches for parent findings given the log category & correlation time window & collects all related docs
      * for them.
      */
-    private void searchFindingsByTimestamp(String detectorType, Map<String, List<CorrelationQuery>> categoryToQueriesMap, List<String> correlationRules) {
+    private void searchFindingsByTimestamp(String detectorType, Map<String, List<CorrelationQuery>> categoryToQueriesMap, Map<String, Long> categoryToTimeWindowMap, List<String> correlationRules) {
         long findingTimestamp = request.getFinding().getTimestamp().toEpochMilli();
         MultiSearchRequest mSearchRequest = new MultiSearchRequest();
         List<Pair<String, List<CorrelationQuery>>> categoryToQueriesPairs = new ArrayList<>();
 
         for (Map.Entry<String, List<CorrelationQuery>> categoryToQueries: categoryToQueriesMap.entrySet()) {
             RangeQueryBuilder queryBuilder = QueryBuilders.rangeQuery("timestamp")
-                    .gte(findingTimestamp - corrTimeWindow)
-                    .lte(findingTimestamp + corrTimeWindow);
+                    .gte(findingTimestamp - categoryToTimeWindowMap.get(categoryToQueries.getKey()))
+                    .lte(findingTimestamp + categoryToTimeWindowMap.get(categoryToQueries.getKey()));
 
             SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
             searchSourceBuilder.query(queryBuilder);
@@ -222,7 +262,7 @@
             searchRequest.indices(DetectorMonitorConfig.getAllFindingsIndicesPattern(categoryToQueries.getKey()));
             searchRequest.source(searchSourceBuilder);
             mSearchRequest.add(searchRequest);
-            categoryToQueriesPairs.add(new Pair<>(categoryToQueries.getKey(), categoryToQueries.getValue()));
+            categoryToQueriesPairs.add(Pair.of(categoryToQueries.getKey(), categoryToQueries.getValue()));
         }
 
         if (!mSearchRequest.requests().isEmpty()) {
@@ -246,17 +286,17 @@
                                     .map(Object::toString).collect(Collectors.toList()));
                         }
 
-                        List<CorrelationQuery> correlationQueries = categoryToQueriesPairs.get(idx).getSecond();
+                        List<CorrelationQuery> correlationQueries = categoryToQueriesPairs.get(idx).getValue();
                         List<String> indices = correlationQueries.stream().map(CorrelationQuery::getIndex).collect(Collectors.toList());
                         List<String> queries = correlationQueries.stream().map(CorrelationQuery::getQuery).collect(Collectors.toList());
-                        relatedDocsMap.put(categoryToQueriesPairs.get(idx).getFirst(),
+                        relatedDocsMap.put(categoryToQueriesPairs.get(idx).getKey(),
                                 new DocSearchCriteria(
                                         indices,
                                         queries,
                                         relatedDocIds));
                         ++idx;
                     }
-                    searchDocsWithFilterKeys(detectorType, relatedDocsMap, correlationRules);
+                    searchDocsWithFilterKeys(detectorType, relatedDocsMap, categoryToTimeWindowMap, correlationRules);
                 }
 
                 @Override
@@ -272,7 +312,7 @@
     /**
      * Given the related docs from parent findings, this method filters only those related docs which match parent join criteria.
      */
-    private void searchDocsWithFilterKeys(String detectorType, Map<String, DocSearchCriteria> relatedDocsMap, List<String> correlationRules) {
+    private void searchDocsWithFilterKeys(String detectorType, Map<String, DocSearchCriteria> relatedDocsMap, Map<String, Long> categoryToTimeWindowMap, List<String> correlationRules) {
         MultiSearchRequest mSearchRequest = new MultiSearchRequest();
         List<String> categories = new ArrayList<>();
 
@@ -283,6 +323,7 @@
             for (String query: docSearchCriteria.getValue().queries) {
                 queryBuilder = queryBuilder.should(QueryBuilders.queryStringQuery(query));
             }
+            queryBuilder.minimumShouldMatch(1).boost(1.0f);
 
             SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
             searchSourceBuilder.query(queryBuilder);
@@ -319,7 +360,7 @@
                         filteredRelatedDocIds.put(categories.get(idx), docIds);
                         ++idx;
                     }
-                    getCorrelatedFindings(detectorType, filteredRelatedDocIds, correlationRules);
+                    getCorrelatedFindings(detectorType, filteredRelatedDocIds, categoryToTimeWindowMap, correlationRules);
                 }
 
                 @Override
@@ -336,16 +377,16 @@
      * Given the filtered related docs of the parent findings, this method gets the actual filtered parent findings for
      * the finding to be correlated.
      */
-    private void getCorrelatedFindings(String detectorType, Map<String, List<String>> filteredRelatedDocIds, List<String> correlationRules) {
+    private void getCorrelatedFindings(String detectorType, Map<String, List<String>> filteredRelatedDocIds, Map<String, Long> categoryToTimeWindowMap, List<String> correlationRules) {
         long findingTimestamp = request.getFinding().getTimestamp().toEpochMilli();
         MultiSearchRequest mSearchRequest = new MultiSearchRequest();
         List<String> categories = new ArrayList<>();
 
         for (Map.Entry<String, List<String>> relatedDocIds: filteredRelatedDocIds.entrySet()) {
             BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery()
                     .filter(QueryBuilders.rangeQuery("timestamp")
-                            .gte(findingTimestamp - corrTimeWindow)
-                            .lte(findingTimestamp + corrTimeWindow))
+                            .gte(findingTimestamp - categoryToTimeWindowMap.get(relatedDocIds.getKey()))
+                            .lte(findingTimestamp + categoryToTimeWindowMap.get(relatedDocIds.getKey())))
                     .must(QueryBuilders.termsQuery("correlated_doc_ids", relatedDocIds.getValue()));
 
             if (relatedDocIds.getKey().equals(detectorType)) {

src/main/java/org/opensearch/securityanalytics/model/CorrelationQuery.java

@@ -22,40 +22,53 @@
     private static final String QUERY = "query";
     private static final String CATEGORY = "category";
 
+    private static final String FIELD = "field";
+
     private String index;
 
     private String query;
 
     private String category;
 
-    public CorrelationQuery(String index, String query, String category) {
+    private String field;
+
+    public CorrelationQuery(String index, String query, String category, String field) {
         this.index = index;
         this.query = query;
         this.category = category;
+        this.field = field;
     }
 
     public CorrelationQuery(StreamInput sin) throws IOException {
-        this(sin.readString(), sin.readString(), sin.readString());
+        this(sin.readString(), sin.readOptionalString(), sin.readString(), sin.readOptionalString());
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(index);
-        out.writeString(query);
+        out.writeOptionalString(query);
         out.writeString(category);
+        out.writeOptionalString(field);
     }
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
-        builder.field(INDEX, index).field(QUERY, query).field(CATEGORY, category);
+        builder.field(INDEX, index).field(CATEGORY, category);
+        if (query != null) {
+            builder.field(QUERY, query);
+        }
+        if (field != null) {
+            builder.field(FIELD, field);
+        }
         return builder.endObject();
     }
 
     public static CorrelationQuery parse(XContentParser xcp) throws IOException {
         String index = null;
         String query = null;
         String category = null;
+        String field = null;
 
         XContentParserUtils.ensureExpectedToken(XContentParser.Token.START_OBJECT, xcp.currentToken(), xcp);
         while (xcp.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -72,11 +85,14 @@
                 case CATEGORY:
                     category = xcp.text();
                     break;
+                case FIELD:
+                    field = xcp.text();
+                    break;
                 default:
                     xcp.skipChildren();
             }
         }
-        return new CorrelationQuery(index, query, category);
+        return new CorrelationQuery(index, query, category, field);
     }
 
     public static CorrelationQuery readFrom(StreamInput sin) throws IOException {
@@ -94,4 +110,8 @@
     public String getCategory() {
         return category;
     }
+
+    public String getField() {
+        return field;
+    }
 }