fixed incorrect total count of entries in final result for list threa…

…t apis
akto-api-security · Jan 9, 2025 · af16337 · af16337
1 parent 20a4cdd
commit af16337
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 16 deletions.
diff --git a/...etection-backend/src/main/java/com/akto/threat/backend/service/MaliciousEventService.java b/...etection-backend/src/main/java/com/akto/threat/backend/service/MaliciousEventService.java
@@ -143,6 +143,7 @@ public ListMaliciousRequestsResponse listMaliciousRequests(
                 MongoDBCollection.ThreatDetection.MALICIOUS_EVENTS, MaliciousEventModel.class);
 
     BasicDBObject query = new BasicDBObject();
+    long total = coll.countDocuments(query);
     try (MongoCursor<MaliciousEventModel> cursor =
         coll.find(query)
             .sort(new BasicDBObject("detectedAt", sort.getOrDefault("detectedAt", -1)))
@@ -169,7 +170,7 @@ public ListMaliciousRequestsResponse listMaliciousRequests(
                 .build());
       }
       return ListMaliciousRequestsResponse.newBuilder()
-          .setTotal(maliciousEvents.size())
+          .setTotal(total)
           .addAllMaliciousEvents(maliciousEvents)
           .build();
     }

diff --git a/...t-detection-backend/src/main/java/com/akto/threat/backend/service/ThreatActorService.java b/...t-detection-backend/src/main/java/com/akto/threat/backend/service/ThreatActorService.java
@@ -31,9 +31,9 @@ public ListThreatActorResponse listThreatActors(
             .getDatabase(accountId)
             .getCollection(MongoDBCollection.ThreatDetection.MALICIOUS_EVENTS, Document.class);
 
-    List<Document> pipeline = new ArrayList<>();
-    pipeline.add(new Document("$sort", new Document("actor", 1).append("detectedAt", -1))); // sort
-    pipeline.add(
+    List<Document> base = new ArrayList<>();
+    base.add(new Document("$sort", new Document("actor", 1).append("detectedAt", -1))); // sort
+    base.add(
         new Document(
             "$group",
             new Document("_id", "$actor")
@@ -42,6 +42,15 @@ public ListThreatActorResponse listThreatActors(
                 .append("latestApiIp", new Document("$last", "$latestApiIp"))
                 .append("country", new Document("$last", "$country"))
                 .append("discoveredAt", new Document("$last", "$detectedAt"))));
+
+    List<Document> countPipeline = new ArrayList<>(base);
+    countPipeline.add(new Document("$count", "total"));
+
+    Document result = coll.aggregate(countPipeline).first();
+    long total = result != null ? result.getInteger("total", 0) : 0;
+
+    List<Document> pipeline = new ArrayList<>(base);
+
     pipeline.add(new Document("$skip", skip));
     pipeline.add(new Document("$limit", limit));
 
@@ -65,10 +74,7 @@ public ListThreatActorResponse listThreatActors(
       }
     }
 
-    return ListThreatActorResponse.newBuilder()
-        .addAllActors(actors)
-        .setTotal(actors.size())
-        .build();
+    return ListThreatActorResponse.newBuilder().addAllActors(actors).setTotal(total).build();
   }
 
   public ThreatActorByCountryResponse getThreatActorByCountry(

diff --git a/...eat-detection-backend/src/main/java/com/akto/threat/backend/service/ThreatApiService.java b/...eat-detection-backend/src/main/java/com/akto/threat/backend/service/ThreatApiService.java
@@ -30,14 +30,14 @@ public ListThreatApiResponse listThreatApis(String accountId, ListThreatApiReque
             .getDatabase(accountId)
             .getCollection(MongoDBCollection.ThreatDetection.MALICIOUS_EVENTS, Document.class);
 
-    List<Document> pipeline = new ArrayList<>();
-    pipeline.add(
+    List<Document> base = new ArrayList<>();
+    base.add(
         new Document(
             "$sort",
             new Document("latestApiEndpoint", 1)
                 .append("latestApiMethod", 1)
                 .append("detectedAt", -1))); // sort
-    pipeline.add(
+    base.add(
         new Document(
             "$group",
             new Document(
@@ -47,9 +47,18 @@ public ListThreatApiResponse listThreatApis(String accountId, ListThreatApiReque
                 .append("discoveredAt", new Document("$last", "$detectedAt"))
                 .append("distinctActors", new Document("$addToSet", "$actor"))
                 .append("requestsCount", new Document("$sum", 1))));
-    pipeline.add(
+    base.add(
         new Document(
             "$addFields", new Document("actorsCount", new Document("$size", "$distinctActors"))));
+
+    List<Document> countPipeline = new ArrayList<>(base);
+    countPipeline.add(new Document("$count", "total"));
+
+    Document result = coll.aggregate(countPipeline).first();
+    long total = result != null ? result.getInteger("total", 0) : 0;
+
+    List<Document> pipeline = new ArrayList<>(base);
+
     pipeline.add(new Document("$project", new Document("distinctActors", 0)));
     pipeline.add(new Document("$skip", skip));
     pipeline.add(new Document("$limit", limit));
@@ -79,7 +88,7 @@ public ListThreatApiResponse listThreatApis(String accountId, ListThreatApiReque
       e.printStackTrace();
     }
 
-    return ListThreatApiResponse.newBuilder().addAllApis(apis).setTotal(apis.size()).build();
+    return ListThreatApiResponse.newBuilder().addAllApis(apis).setTotal(total).build();
   }
 
   public ThreatCategoryWiseCountResponse getSubCategoryWiseCount(

diff --git a/protobuf/threat_detection/service/dashboard_service/v1/service.proto b/protobuf/threat_detection/service/dashboard_service/v1/service.proto
@@ -25,7 +25,7 @@ message ListMaliciousRequestsResponse {
       string sub_category = 13;
     }
     repeated MaliciousEvent malicious_events = 1;
-    int32 total = 2;
+    uint64 total = 2;
 }
 
 message ListMaliciousRequestsRequest {
@@ -58,7 +58,7 @@ message ListThreatActorResponse {
     string country = 6;
   }
   repeated ThreatActor actors = 1;
-  int32 total = 2;
+  uint64 total = 2;
 }
 
 message ListThreatApiRequest {
@@ -77,7 +77,7 @@ message ListThreatApiResponse {
   }
 
   repeated ThreatApi apis = 1;
-  int32 total = 2;
+  uint64 total = 2;
 }
 
 message ThreatActorByCountryRequest {