Merge pull request #734 from akto-api-security/hotfix/sample_msg_for_…

…collections Hotfix/sample msg for collections
akto-api-security · Jun 29, 2023 · 1552991 · 1552991
2 parents 577aee8 + 30c4ca0
commit 1552991
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 11 deletions.
diff --git a/apps/testing/src/main/java/com/akto/store/SampleMessageStore.java b/apps/testing/src/main/java/com/akto/store/SampleMessageStore.java
@@ -18,6 +18,7 @@
 import com.akto.testing.NucleiExecutor;
 import com.mongodb.BasicDBObject;
 import com.mongodb.client.model.Filters;
+import org.bson.conversions.Bson;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -76,8 +77,9 @@ public static List<TestRoles> fetchTestRoles() {
     }
 
 
-    public static Map<ApiInfo.ApiInfoKey, List<String>> fetchSampleMessages() {
-        List<SampleData> sampleDataList = SampleDataDao.instance.findAll(new BasicDBObject(), 0, 10_000, null);
+    public static Map<ApiInfo.ApiInfoKey, List<String>> fetchSampleMessages(Set<Integer> apiCollectionIds) {
+        Bson filterQ = Filters.in("_id.apiCollectionId", apiCollectionIds);
+        List<SampleData> sampleDataList = SampleDataDao.instance.findAll(filterQ, 0, 10_000, null);
         System.out.println("SampleDataSize " + sampleDataList.size());
         Map<ApiInfo.ApiInfoKey, List<String>> tempSampleDataMap = new HashMap<>();
         for (SampleData sampleData: sampleDataList) {

diff --git a/apps/testing/src/main/java/com/akto/testing/TestExecutor.java b/apps/testing/src/main/java/com/akto/testing/TestExecutor.java
@@ -95,14 +95,31 @@ public void workflowInit (TestingRun testingRun, ObjectId summaryId) {
         );
     }
 
+    private Set<Integer> extractApiCollectionIds(List<ApiInfo.ApiInfoKey> apiInfoKeyList) {
+        Set<Integer> ret = new HashSet<>();
+        for(ApiInfo.ApiInfoKey apiInfoKey: apiInfoKeyList) {
+            ret.add(apiInfoKey.getApiCollectionId());
+        }
+
+        return ret;
+    }
+
     public void apiWiseInit(TestingRun testingRun, ObjectId summaryId) {
         int accountId = Context.accountId.get();
         int now = Context.now();
         int maxConcurrentRequests = testingRun.getMaxConcurrentRequests() > 0 ? testingRun.getMaxConcurrentRequests() : 100;
         TestingEndpoints testingEndpoints = testingRun.getTestingEndpoints();
 
         Map<String, SingleTypeInfo> singleTypeInfoMap = SampleMessageStore.buildSingleTypeInfoMap(testingEndpoints);
-        Map<ApiInfo.ApiInfoKey, List<String>> sampleMessages = SampleMessageStore.fetchSampleMessages();
+
+        List<ApiInfo.ApiInfoKey> apiInfoKeyList = testingEndpoints.returnApis();
+        if (apiInfoKeyList == null || apiInfoKeyList.isEmpty()) return;
+        loggerMaker.infoAndAddToDb("APIs found: " + apiInfoKeyList.size(), LogDb.TESTING);
+
+
+        Set<Integer> apiCollectionIds = extractApiCollectionIds(apiInfoKeyList);
+
+        Map<ApiInfo.ApiInfoKey, List<String>> sampleMessages = SampleMessageStore.fetchSampleMessages(apiCollectionIds);
         List<TestRoles> testRoles = SampleMessageStore.fetchTestRoles();
         AuthMechanism authMechanism = AuthMechanismsDao.instance.findOne(new BasicDBObject());
 
@@ -125,10 +142,6 @@ public void apiWiseInit(TestingRun testingRun, ObjectId summaryId) {
             return;
         }
 
-        List<ApiInfo.ApiInfoKey> apiInfoKeyList = testingEndpoints.returnApis();
-        if (apiInfoKeyList == null || apiInfoKeyList.isEmpty()) return;
-        loggerMaker.infoAndAddToDb("APIs found: " + apiInfoKeyList.size(), LogDb.TESTING);
-
         Map<ApiInfo.ApiInfoKey, List<String>> sampleDataMapForStatusCodeAnalyser = new HashMap<>();
         Set<ApiInfo.ApiInfoKey> apiInfoKeySet = new HashSet<>(apiInfoKeyList);
         for (ApiInfo.ApiInfoKey apiInfoKey: sampleMessages.keySet()) {

diff --git a/apps/testing/src/test/java/com/akto/store/TestSampleMessageStore.java b/apps/testing/src/test/java/com/akto/store/TestSampleMessageStore.java
@@ -23,19 +23,27 @@ public void testFetchSampleMessages() {
         SampleData sampleData1 = new SampleData(new Key(0, "url1", URLMethods.Method.GET,0,0,0), null);
         SampleData sampleData2 = new SampleData(new Key(0, "url2", URLMethods.Method.GET,0,0,0), Arrays.asList("m1", "m2"));
         SampleData sampleData3 = new SampleData(new Key(0, "url3", URLMethods.Method.GET,0,0,0), Collections.emptyList());
-        SampleDataDao.instance.insertMany(Arrays.asList(sampleData1, sampleData2, sampleData3));
+        SampleData sampleData4 = new SampleData(new Key(1, "url1", URLMethods.Method.GET,0,0,0), Arrays.asList("m3", "m4", "m5"));
+        SampleDataDao.instance.insertMany(Arrays.asList(sampleData1, sampleData2, sampleData3, sampleData4));
 
-        Map<ApiInfo.ApiInfoKey, List<String>> sampleDataMap =  SampleMessageStore.fetchSampleMessages();
+        Set<Integer> apiCollectionIds = new HashSet<>();
+        apiCollectionIds.add(0);
+        apiCollectionIds.add(1);
 
-        assertEquals(sampleDataMap.size(), 2);
+        Map<ApiInfo.ApiInfoKey, List<String>> sampleDataMap =  SampleMessageStore.fetchSampleMessages(apiCollectionIds);
+
+        assertEquals(sampleDataMap.size(), 3);
         List<String> messages = sampleDataMap.get(new ApiInfo.ApiInfoKey(0, "url2", URLMethods.Method.GET));
         assertEquals(messages.size(), 2);
 
+        messages = sampleDataMap.get(new ApiInfo.ApiInfoKey(1, "url1", URLMethods.Method.GET));
+        assertEquals(messages.size(), 3);
+
         SampleDataDao.instance.getMCollection().drop();
         sampleData2 = new SampleData(new Key(0, "url2", URLMethods.Method.GET,0,0,0), Arrays.asList("m1", "m2", "m3"));
         SampleDataDao.instance.insertMany(Arrays.asList(sampleData1, sampleData2));
 
-        sampleDataMap =  SampleMessageStore.fetchSampleMessages();
+        sampleDataMap =  SampleMessageStore.fetchSampleMessages(apiCollectionIds);
         assertEquals(sampleDataMap.size(), 1);
         messages = sampleDataMap.get(new ApiInfo.ApiInfoKey(0, "url2", URLMethods.Method.GET));
         assertEquals(messages.size(), 3);