Merge pull request #730 from akto-api-security/feature/pii_cleaner_wh…

…itelist Feature/pii cleaner whitelist
akto-api-security · Jun 26, 2023 · 67d6e2c · 67d6e2c
2 parents 963e97c + ffae34c
commit 67d6e2c
Showing 1 changed file with 21 additions and 16 deletions.
diff --git a/apps/dashboard/src/main/java/com/akto/listener/InitializerListener.java b/apps/dashboard/src/main/java/com/akto/listener/InitializerListener.java
@@ -112,13 +112,16 @@ public static String getDomain() {
     }
 
     public void setUpPiiCleanerScheduler(){
+        Set<Integer> whiteListCollectionSet = new HashSet<>();
+        whiteListCollectionSet.add(-122281555);
+
         scheduler.scheduleAtFixedRate(new Runnable() {
             public void run() {
                 String mongoURI = System.getenv("AKTO_MONGO_CONN");
                 DaoInit.init(new ConnectionString(mongoURI));
                 Context.accountId.set(1_000_000);
                 try {
-                    executePiiCleaner(true);
+                    executePiiCleaner(whiteListCollectionSet);
                 } catch (Exception e) {
 
                 }
@@ -273,7 +276,7 @@ static void executeTestSourcesFetch() {
 
     }
 
-    static void executePiiCleaner(boolean isDryRun) {
+    static void executePiiCleaner(Set<Integer> whiteListCollectionSet) {
         final int BATCH_SIZE = 100;
         int currMarker = 0;
         Bson filterSsdQ =
@@ -312,18 +315,21 @@ static void executePiiCleaner(boolean isDryRun) {
                 dataPoints++;
             }
 
-            bulkSensitiveInvalidate(idsToDelete, isDryRun);
-            bulkSingleTypeInfoDelete(idsToDelete, isDryRun);
+            bulkSensitiveInvalidate(idsToDelete, whiteListCollectionSet);
+            bulkSingleTypeInfoDelete(idsToDelete, whiteListCollectionSet);
 
         } while (dataPoints == BATCH_SIZE);
     }
 
-    private static void bulkSensitiveInvalidate(List<SingleTypeInfo.ParamId> idsToDelete, boolean isDryRun) {
+    private static void bulkSensitiveInvalidate(List<SingleTypeInfo.ParamId> idsToDelete, Set<Integer> whiteListCollectionSet) {
         ArrayList<WriteModel<SensitiveSampleData>> bulkSensitiveInvalidateUpdates = new ArrayList<>();
         for(SingleTypeInfo.ParamId paramId: idsToDelete) {
             String paramStr = "PII cleaner - invalidating: " + paramId.getApiCollectionId() + ": " + paramId.getMethod() + " " + paramId.getUrl() + " > " + paramId.getParam();
             String url = "dashboard/observe/inventory/"+paramId.getApiCollectionId()+"/"+Base64.getEncoder().encodeToString((paramId.getUrl() + " " + paramId.getMethod()).getBytes());
             loggerMaker.infoAndAddToDb(paramStr + url, LogDb.DASHBOARD);
+
+            if (!whiteListCollectionSet.contains(paramId.getApiCollectionId())) continue;
+
             List<Bson> filters = new ArrayList<>();
             filters.add(Filters.eq("url", paramId.getUrl()));
             filters.add(Filters.eq("method", paramId.getMethod()));
@@ -336,21 +342,22 @@ private static void bulkSensitiveInvalidate(List<SingleTypeInfo.ParamId> idsToDe
         }
 
         if (!bulkSensitiveInvalidateUpdates.isEmpty()) {
-            if (!isDryRun) {
-                BulkWriteResult bwr =
-                        SensitiveSampleDataDao.instance.getMCollection().bulkWrite(bulkSensitiveInvalidateUpdates, new BulkWriteOptions().ordered(false));
+            BulkWriteResult bwr =
+                    SensitiveSampleDataDao.instance.getMCollection().bulkWrite(bulkSensitiveInvalidateUpdates, new BulkWriteOptions().ordered(false));
 
-                loggerMaker.infoAndAddToDb("PII cleaner - modified " + bwr.getModifiedCount() + " from STI", LogDb.DASHBOARD);
-            }
+            loggerMaker.infoAndAddToDb("PII cleaner - modified " + bwr.getModifiedCount() + " from STI", LogDb.DASHBOARD);
         }
 
     }
 
-    private static void bulkSingleTypeInfoDelete(List<SingleTypeInfo.ParamId> idsToDelete, boolean isDryRun) {
+    private static void bulkSingleTypeInfoDelete(List<SingleTypeInfo.ParamId> idsToDelete, Set<Integer> whiteListCollectionSet) {
         ArrayList<WriteModel<SingleTypeInfo>> bulkUpdatesForSingleTypeInfo = new ArrayList<>();
         for(SingleTypeInfo.ParamId paramId: idsToDelete) {
             String paramStr = "PII cleaner - deleting: " + paramId.getApiCollectionId() + ": " + paramId.getMethod() + " " + paramId.getUrl() + " > " + paramId.getParam();
             loggerMaker.infoAndAddToDb(paramStr, LogDb.DASHBOARD);
+
+            if (!whiteListCollectionSet.contains(paramId.getApiCollectionId())) continue;
+
             List<Bson> filters = new ArrayList<>();
             filters.add(Filters.eq("url", paramId.getUrl()));
             filters.add(Filters.eq("method", paramId.getMethod()));
@@ -363,12 +370,10 @@ private static void bulkSingleTypeInfoDelete(List<SingleTypeInfo.ParamId> idsToD
         }
 
         if (!bulkUpdatesForSingleTypeInfo.isEmpty()) {
-            if (!isDryRun) {
-                BulkWriteResult bwr =
-                        SingleTypeInfoDao.instance.getMCollection().bulkWrite(bulkUpdatesForSingleTypeInfo, new BulkWriteOptions().ordered(false));
+            BulkWriteResult bwr =
+                    SingleTypeInfoDao.instance.getMCollection().bulkWrite(bulkUpdatesForSingleTypeInfo, new BulkWriteOptions().ordered(false));
 
-                loggerMaker.infoAndAddToDb("PII cleaner - deleted " + bwr.getDeletedCount() + " from STI", LogDb.DASHBOARD);
-            }
+            loggerMaker.infoAndAddToDb("PII cleaner - deleted " + bwr.getDeletedCount() + " from STI", LogDb.DASHBOARD);
         }
 
     }