Skip to content

Commit a6e56e4

Browse files
ankush-jain-aktoankush-jain-akto
authored
Merge pull request #2608 from akto-api-security/threat_query_optimizations
optimise fetch threat actor api
2 parents 245aa94 + 5e4f9c6 commit a6e56e4

File tree

1 file changed

+79
-101
lines changed

1 file changed

+79
-101
lines changed

‎apps/threat-detection-backend/src/main/java/com/akto/threat/backend/service/ThreatActorService.java

Lines changed: 79 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333

3434
import java.util.ArrayList;
3535
import java.util.Arrays;
36+
import java.util.Collections;
3637
import java.util.List;
3738
import java.util.Map;
3839

@@ -112,122 +113,99 @@ public ThreatConfiguration modifyThreatConfiguration(String accountId, ThreatCon
112113
return builder.build();
113114
}
114115

115-
public ListThreatActorResponse listThreatActors(
116-
String accountId, ListThreatActorsRequest request) {
117-
118-
int skip = request.hasSkip() ? request.getSkip() : 0;
119-
int limit = request.getLimit();
120-
Map<String, Integer> sort = request.getSortMap();
121-
MongoCollection<Document> coll =
122-
this.mongoClient
123-
.getDatabase(accountId)
124-
.getCollection(MongoDBCollection.ThreatDetection.MALICIOUS_EVENTS, Document.class);
125-
126-
ListThreatActorsRequest.Filter filter = request.getFilter();
127-
128-
List<Document> base = new ArrayList<>();
129-
130-
Document match = new Document();
131-
132-
if (!filter.getActorsList().isEmpty()) {
133-
match.append("actor", new Document("$in", filter.getActorsList()));
134-
}
135-
136-
if (!filter.getLatestIpsList().isEmpty()) {
137-
match.append("latestApiIp", new Document("$in", filter.getLatestIpsList()));
138-
}
139-
140-
if (!filter.getLatestAttackList().isEmpty()) {
141-
match.append("subCategory", new Document("$in", filter.getLatestAttackList()));
142-
}
143-
144-
if (!filter.getCountryList().isEmpty()) {
145-
match.append("country", new Document("$in", filter.getCountryList()));
146-
}
147-
148-
if (filter.hasDetectedAtTimeRange()) {
149-
long start = filter.getDetectedAtTimeRange().getStart();
150-
long end = filter.getDetectedAtTimeRange().getEnd();
151-
match.append("detectedAt", new Document("$gte", start).append("$lte", end));
152-
}
153-
154-
if (request.getStartTs() != 0 && request.getEndTs() != 0) {
155-
long start = request.getStartTs();
156-
long end = request.getEndTs();
157-
match.append("detectedAt", new Document("$gte", start).append("$lte", end));
158-
}
159116

160-
if (!match.isEmpty()) {
161-
base.add(new Document("$match", match));
162-
}
163-
164-
base.add(new Document("$sort", new Document("detectedAt", -1)));
165-
base.add(new Document("$skip", skip));
166-
base.add(new Document("$limit", limit));
167-
base.add(
168-
new Document(
169-
"$group",
170-
new Document("_id", "$actor")
171-
.append("latestApiEndpoint", new Document("$last", "$latestApiEndpoint"))
172-
.append("latestApiMethod", new Document("$last", "$latestApiMethod"))
173-
.append("latestApiIp", new Document("$last", "$latestApiIp"))
174-
.append("country", new Document("$last", "$country"))
175-
.append("discoveredAt", new Document("$last", "$detectedAt"))
176-
.append("latestSubCategory", new Document("$last", "$subCategory"))));
177-
178-
List<Document> countPipeline = new ArrayList<>(base);
179-
countPipeline.add(new Document("$count", "total"));
117+
public ListThreatActorResponse listThreatActors(String accountId, ListThreatActorsRequest request) {
118+
int skip = request.hasSkip() ? request.getSkip() : 0;
119+
int limit = request.getLimit();
120+
Map<String, Integer> sort = request.getSortMap();
180121

181-
Document result = coll.aggregate(countPipeline).first();
182-
long total = result != null ? result.getInteger("total", 0) : 0;
122+
MongoCollection<Document> coll = this.mongoClient
123+
.getDatabase(accountId)
124+
.getCollection(MongoDBCollection.ThreatDetection.MALICIOUS_EVENTS, Document.class);
183125

184-
List<Document> pipeline = new ArrayList<>(base);
126+
ListThreatActorsRequest.Filter filter = request.getFilter();
127+
Document match = new Document();
185128

186-
pipeline.add(new Document("$skip", skip));
187-
pipeline.add(new Document("$limit", limit));
129+
// Apply filters
130+
if (!filter.getActorsList().isEmpty()) match.append("actor", new Document("$in", filter.getActorsList()));
131+
if (!filter.getLatestIpsList().isEmpty()) match.append("latestApiIp", new Document("$in", filter.getLatestIpsList()));
132+
if (!filter.getLatestAttackList().isEmpty()) match.append("subCategory", new Document("$in", filter.getLatestAttackList()));
133+
if (!filter.getCountryList().isEmpty()) match.append("country", new Document("$in", filter.getCountryList()));
134+
if (filter.hasDetectedAtTimeRange()) {
135+
match.append("detectedAt", new Document("$gte", filter.getDetectedAtTimeRange().getStart()).append("$lte", filter.getDetectedAtTimeRange().getEnd()));
136+
}
137+
if (request.getStartTs() != 0 && request.getEndTs() != 0) {
138+
match.append("detectedAt", new Document("$gte", request.getStartTs()).append("$lte", request.getEndTs()));
139+
}
188140

189-
pipeline.add(
190-
new Document(
191-
"$sort", new Document("discoveredAt", sort.getOrDefault("discoveredAt", -1)))); // sort
141+
List<Document> pipeline = new ArrayList<>();
142+
if (!match.isEmpty()) pipeline.add(new Document("$match", match));
143+
144+
// Sort first for $first to work
145+
pipeline.add(new Document("$sort", new Document("detectedAt", -1)));
146+
147+
pipeline.add(new Document("$group", new Document("_id", "$actor")
148+
.append("latestApiEndpoint", new Document("$first", "$latestApiEndpoint"))
149+
.append("latestApiMethod", new Document("$first", "$latestApiMethod"))
150+
.append("latestApiIp", new Document("$first", "$latestApiIp"))
151+
.append("country", new Document("$first", "$country"))
152+
.append("discoveredAt", new Document("$first", "$detectedAt"))
153+
.append("latestSubCategory", new Document("$first", "$subCategory"))
154+
));
155+
156+
// Facet: count and paginated result
157+
List<Document> facetStages = Arrays.asList(
158+
new Document("$sort", new Document("discoveredAt", sort.getOrDefault("discoveredAt", -1))),
159+
new Document("$skip", skip),
160+
new Document("$limit", limit)
161+
);
192162

193-
List<ListThreatActorResponse.ThreatActor> actors = new ArrayList<>();
194-
try (MongoCursor<Document> cursor = coll.aggregate(pipeline).cursor()) {
195-
while (cursor.hasNext()) {
196-
Document doc = cursor.next();
197-
198-
Bson filters = Filters.eq("actor", doc.getString("_id"));
199-
Bson sort2 = Sorts.descending("detectedAt");
200-
MongoCursor<Document> cursor2 = coll.find(filters).sort(sort2).limit(40).cursor();
201-
List<ActivityData> activityDataList = new ArrayList<>();
202-
while (cursor2.hasNext()) {
203-
Document doc2 = cursor2.next();
204-
activityDataList.add(
205-
ActivityData.newBuilder()
206-
.setUrl(doc2.getString("latestApiEndpoint"))
207-
.setDetectedAt(doc2.getLong("detectedAt"))
208-
.setSubCategory(doc2.getString("subCategory"))
209-
.setSeverity(doc2.getString("severity"))
210-
.setMethod(doc2.getString("latestApiMethod"))
211-
.build()
212-
);
213-
}
163+
pipeline.add(new Document("$facet", new Document()
164+
.append("paginated", facetStages)
165+
.append("count", Arrays.asList(new Document("$count", "total")))
166+
));
167+
168+
Document result = coll.aggregate(pipeline).first();
169+
List<Document> paginated = result.getList("paginated", Document.class, Collections.emptyList());
170+
List<Document> countList = result.getList("count", Document.class, Collections.emptyList());
171+
long total = countList.isEmpty() ? 0 : countList.get(0).getInteger("total");
172+
173+
// Activity fetch
174+
List<ListThreatActorResponse.ThreatActor> actors = new ArrayList<>();
175+
for (Document doc : paginated) {
176+
String actorId = doc.getString("_id");
177+
List<ActivityData> activityDataList = new ArrayList<>();
178+
179+
try (MongoCursor<Document> cursor2 = coll.find(Filters.eq("actor", actorId))
180+
.sort(Sorts.descending("detectedAt"))
181+
.limit(40)
182+
.cursor()) {
183+
while (cursor2.hasNext()) {
184+
Document doc2 = cursor2.next();
185+
activityDataList.add(ActivityData.newBuilder()
186+
.setUrl(doc2.getString("latestApiEndpoint"))
187+
.setDetectedAt(doc2.getLong("detectedAt"))
188+
.setSubCategory(doc2.getString("subCategory"))
189+
.setSeverity(doc2.getString("severity"))
190+
.setMethod(doc2.getString("latestApiMethod"))
191+
.build());
192+
}
193+
}
214194

215-
actors.add(
216-
ListThreatActorResponse.ThreatActor.newBuilder()
217-
.setId(doc.getString("_id"))
195+
actors.add(ListThreatActorResponse.ThreatActor.newBuilder()
196+
.setId(actorId)
218197
.setLatestApiEndpoint(doc.getString("latestApiEndpoint"))
219198
.setLatestApiMethod(doc.getString("latestApiMethod"))
220199
.setLatestApiIp(doc.getString("latestApiIp"))
221200
.setDiscoveredAt(doc.getLong("discoveredAt"))
222201
.setCountry(doc.getString("country"))
223-
.addAllActivityData(activityDataList)
224202
.setLatestSubcategory(doc.getString("latestSubCategory"))
203+
.addAllActivityData(activityDataList)
225204
.build());
226-
}
227-
}
205+
}
228206

229-
return ListThreatActorResponse.newBuilder().addAllActors(actors).setTotal(total).build();
230-
}
207+
return ListThreatActorResponse.newBuilder().addAllActors(actors).setTotal(total).build();
208+
}
231209

232210
public DailyActorsCountResponse getDailyActorCounts(String accountId, long startTs, long endTs) {
233211

0 commit comments

Comments
 (0)