enable multiple attacker tokens

akto-api-security · Nov 5, 2024 · e0395ce · e0395ce
1 parent 8e598f5
commit e0395ce
Show file tree

Hide file tree

Showing 9 changed files with 70 additions and 14 deletions.
diff --git a/apps/dashboard/src/main/java/com/akto/action/test_editor/SaveTestEditorAction.java b/apps/dashboard/src/main/java/com/akto/action/test_editor/SaveTestEditorAction.java
@@ -305,7 +305,7 @@ public String runTestForGivenTemplate() {
                 apiInfoKey.getString(ApiInfo.ApiInfoKey.URL),
                 URLMethods.Method.valueOf(apiInfoKey.getString(ApiInfo.ApiInfoKey.METHOD)));
 
-        AuthMechanism authMechanism = TestRolesDao.instance.fetchAttackerToken(0);
+        AuthMechanism authMechanism = TestRolesDao.instance.fetchAttackerToken(0, null);
         Map<ApiInfo.ApiInfoKey, List<String>> sampleDataMap = new HashMap<>();
         Map<ApiInfo.ApiInfoKey, List<String>> newSampleDataMap = new HashMap<>();
 

diff --git a/apps/dashboard/src/main/java/com/akto/action/testing/AuthMechanismAction.java b/apps/dashboard/src/main/java/com/akto/action/testing/AuthMechanismAction.java
@@ -141,7 +141,7 @@ public String triggerSingleLoginFlowStep() {
 
     public String fetchAuthMechanismData() {
 
-        authMechanism = TestRolesDao.instance.fetchAttackerToken(0);
+        authMechanism = TestRolesDao.instance.fetchAttackerToken(0, null);
         return SUCCESS.toUpperCase();
     }
 

diff --git a/apps/dashboard/src/main/java/com/akto/action/testing/StartTestAction.java b/apps/dashboard/src/main/java/com/akto/action/testing/StartTestAction.java
@@ -120,7 +120,7 @@ private TestingRun createTestingRun(int scheduleTimestamp, int periodInSeconds)
             }
         }
 
-        AuthMechanism authMechanism = TestRolesDao.instance.fetchAttackerToken(0);
+        AuthMechanism authMechanism = TestRolesDao.instance.fetchAttackerToken(0, null);
         if (authMechanism == null && testIdConfig == 0) {
             addActionError("Please set authentication mechanism before you test any APIs");
             return null;
@@ -395,7 +395,7 @@ private ArrayList<Bson> getTableFilters(){
 
     public String retrieveAllCollectionTests() {
 
-        this.authMechanism = TestRolesDao.instance.fetchAttackerToken(0);
+        this.authMechanism = TestRolesDao.instance.fetchAttackerToken(0, null);
 
         ArrayList<Bson> testingRunFilters = new ArrayList<>();
         Bson testingRunTypeFilter = getTestingRunTypeFilter(testingRunType);

diff --git a/apps/testing/src/main/java/com/akto/store/AuthMechanismStore.java b/apps/testing/src/main/java/com/akto/store/AuthMechanismStore.java
@@ -2,6 +2,7 @@
 
 import com.akto.dao.AuthMechanismsDao;
 import com.akto.dao.testing.TestRolesDao;
+import com.akto.dto.RawApi;
 import com.akto.dto.testing.AuthMechanism;
 import com.mongodb.BasicDBObject;
 
@@ -10,9 +11,9 @@ public class AuthMechanismStore {
 
     private AuthMechanismStore() {}
 
-    public static AuthMechanismStore create() {
+    public static AuthMechanismStore create(RawApi rawApi) {
         AuthMechanismStore ret = new AuthMechanismStore();
-        ret.authMechanism = TestRolesDao.instance.fetchAttackerToken(0);
+        ret.authMechanism = TestRolesDao.instance.fetchAttackerToken(0, rawApi);
         return ret;
     }
 

diff --git a/apps/testing/src/main/java/com/akto/store/SampleMessageStore.java b/apps/testing/src/main/java/com/akto/store/SampleMessageStore.java
@@ -136,6 +136,22 @@ public static List<RawApi> filterMessagesWithAuthToken(List<RawApi> messages, Au
         return filteredMessages;
     }
 
+    public List<RawApi> findSampleMessages(int k) {
+        List<RawApi> samples = new ArrayList<>();
+        if (sampleDataMap == null) return samples;
+
+        for (ApiInfoKey apiInfoKey : sampleDataMap.keySet()) {
+            List<String> messages = sampleDataMap.getOrDefault(apiInfoKey, new ArrayList<>());
+            if (!messages.isEmpty()) {
+                RawApi rawApi = RawApi.buildFromMessage(messages.get(0));
+                samples.add(rawApi);
+            }
+            if (samples.size() >= k) break;
+        }
+
+        return samples;
+    }
+
     public Map<String, SingleTypeInfo> getSingleTypeInfos() {
         return this.singleTypeInfos;
     }

diff --git a/apps/testing/src/main/java/com/akto/testing/AccessMatrixAnalyzer.java b/apps/testing/src/main/java/com/akto/testing/AccessMatrixAnalyzer.java
@@ -97,8 +97,11 @@ public void run() throws Exception {
             loggerMaker.infoAndAddToDb("Role found: " + roleFromTask, LogDb.TESTING);
             List<TestRoles> testRoles = TestRolesDao.instance.findAll(TestRoles.NAME, roleFromTask);
 
-            AuthMechanismStore authMechanismStore = AuthMechanismStore.create();
+            List<RawApi> rawApis = sampleMessageStore.findSampleMessages(1);
+            RawApi randomRawApi = !rawApis.isEmpty() ? rawApis.get(0) : null;
+            AuthMechanismStore authMechanismStore = AuthMechanismStore.create(randomRawApi);
             AuthMechanism authMechanism = authMechanismStore.getAuthMechanism();
+
             List<CustomAuthType> customAuthTypes = CustomAuthTypeDao.instance.findAll(CustomAuthType.ACTIVE,true);
             TestingUtil testingUtil = new TestingUtil(authMechanism,sampleMessageStore, testRoles,"", customAuthTypes);
 

diff --git a/apps/testing/src/main/java/com/akto/testing/TestExecutor.java b/apps/testing/src/main/java/com/akto/testing/TestExecutor.java
@@ -141,7 +141,10 @@ public void apiWiseInit(TestingRun testingRun, ObjectId summaryId, boolean debug
 
         SampleMessageStore sampleMessageStore = SampleMessageStore.create();
         sampleMessageStore.fetchSampleMessages(Main.extractApiCollectionIds(testingRun.getTestingEndpoints().returnApis()));
-        AuthMechanismStore authMechanismStore = AuthMechanismStore.create();
+
+        List<RawApi> rawApis = sampleMessageStore.findSampleMessages(1);
+        RawApi randomRawApi = !rawApis.isEmpty() ? rawApis.get(0) : null;
+        AuthMechanismStore authMechanismStore = AuthMechanismStore.create(randomRawApi);
 
         List<ApiInfo.ApiInfoKey> apiInfoKeyList = testingEndpoints.returnApis();
         if (apiInfoKeyList == null || apiInfoKeyList.isEmpty()) return;

diff --git a/apps/testing/src/test/java/com/akto/testing/TestExecutorTest.java b/apps/testing/src/test/java/com/akto/testing/TestExecutorTest.java
@@ -67,7 +67,7 @@ private void testFindHostUtil(String url, String answer, String hostName) throws
         Set<Integer> apiCollectionSet = new HashSet<>();
         apiCollectionSet.add(0);
         messageStore.fetchSampleMessages(apiCollectionSet);
-        AuthMechanismStore authMechanismStore = AuthMechanismStore.create();
+        AuthMechanismStore authMechanismStore = AuthMechanismStore.create(null);
         TestingUtil testingUtil = new TestingUtil(authMechanismStore.getAuthMechanism(), messageStore, new ArrayList<>(), "", new ArrayList<>());
 
         String host = TestExecutor.findHost(apiInfoKey, testingUtil.getSampleMessages(), messageStore);

diff --git a/libs/dao/src/main/java/com/akto/dao/testing/TestRolesDao.java b/libs/dao/src/main/java/com/akto/dao/testing/TestRolesDao.java
@@ -4,6 +4,7 @@
 import com.akto.dao.AuthMechanismsDao;
 import com.akto.dao.MCollection;
 import com.akto.dao.context.Context;
+import com.akto.dto.RawApi;
 import com.akto.dto.SensitiveSampleData;
 import com.akto.dto.testing.AuthMechanism;
 import com.akto.dto.testing.TestRoles;
@@ -19,6 +20,8 @@
 import org.bson.types.ObjectId;
 
 import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
 
 public class TestRolesDao extends AccountsContextDao<TestRoles> {
     @Override
@@ -59,14 +62,44 @@ public TestRoles createTestRole (String roleName, ObjectId endpointLogicalGroupI
         return role;
     }
 
-    public AuthMechanism fetchAttackerToken(int apiCollectionId) {
+    public AuthMechanism fetchAttackerToken(int apiCollectionId, RawApi rawApi) {
         TestRoles testRoles = TestRolesDao.instance.findOne(TestRoles.NAME, "ATTACKER_TOKEN_ALL");
         if (testRoles != null && testRoles.getAuthWithCondList().size() > 0) {
-            return testRoles.getAuthWithCondList().get(0).getAuthMechanism();
-        } else {
-            // return AuthMechanismsDao.instance.findOne(new BasicDBObject());
-            return null;
+            List<AuthWithCond> authWithCondList = testRoles.getAuthWithCondList();
+            AuthMechanism defaultAuthMechanism = authWithCondList.get(0).getAuthMechanism();
+            if (rawApi == null) {
+                return defaultAuthMechanism;
+            } else {
+                try {
+                    Map<String, List<String>> reqHeaders = rawApi.getRequest().getHeaders();
+                    for (AuthWithCond authWithCond: authWithCondList)  {
+                        Map<String, String> headerKVPairs = authWithCond.getHeaderKVPairs();
+                        if (headerKVPairs == null) continue;
+
+                        boolean allHeadersMatched = true;
+                        for(String hKey: headerKVPairs.keySet()) {
+                            String hVal = authWithCond.getHeaderKVPairs().get(hKey);
+                            if (reqHeaders.containsKey(hKey.toLowerCase())) {
+                                if (!reqHeaders.get(hKey.toLowerCase()).contains(hVal)) {
+                                    allHeadersMatched = false;
+                                    break;
+                                }
+                            }
+                        }
+
+                        if (allHeadersMatched) {
+                            return authWithCond.getAuthMechanism();
+                        }
+                    }
+                } catch (Exception e) {
+                    return defaultAuthMechanism;
+                }
+            }
+
+            return defaultAuthMechanism;
         }
+
+        return null;
     }
 
     public BasicDBObject fetchAttackerTokenDoc(int apiCollectionId) {