Fixing replace body in test roles

akto-api-security · Nov 12, 2024 · 7951c7f · 7951c7f
1 parent 13e2e61
commit 7951c7f
Show file tree

Hide file tree

Showing 8 changed files with 116 additions and 107 deletions.
diff --git a/libs/dao/src/main/java/com/akto/dto/RawApi.java b/libs/dao/src/main/java/com/akto/dto/RawApi.java
@@ -44,7 +44,7 @@ public BasicDBObject fetchReqPayload() {
             reqBody = "{\"json\": "+reqBody+"}";
         }
         try {
-             payload = BasicDBObject.parse(reqBody);
+            payload = BasicDBObject.parse(reqBody);
         } catch (Exception e) {
             payload = new BasicDBObject();
         }

diff --git a/libs/dao/src/main/java/com/akto/dto/test_editor/Util.java b/libs/dao/src/main/java/com/akto/dto/test_editor/Util.java
@@ -0,0 +1,62 @@
+package com.akto.dto.test_editor;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+
+import com.mongodb.BasicDBList;
+import com.mongodb.BasicDBObject;
+
+public class Util {
+    public static boolean modifyValueInPayload(Object obj, String parentKey, String queryKey, Object queryVal){
+        boolean res = false;
+        if (obj instanceof BasicDBObject) {
+            BasicDBObject basicDBObject = (BasicDBObject) obj;
+
+            Set<String> keySet = basicDBObject.keySet();
+
+            for(String key: keySet) {
+                if (key == null) {
+                    continue;
+                }
+                Object value = basicDBObject.get(key);
+
+                if (!( (value instanceof BasicDBObject) || (value instanceof BasicDBList) )) {
+                    if (key.equalsIgnoreCase(queryKey)) {
+                        basicDBObject.remove(key);
+                        basicDBObject.put(queryKey, queryVal);
+                        return true;
+                    }
+                }
+
+                if (value instanceof BasicDBList) {
+                    BasicDBList valList = (BasicDBList) value;
+                    if (valList.size() == 0 && key.equalsIgnoreCase(queryKey)) {
+                        List<Object> queryList = Collections.singletonList(queryVal);
+                        basicDBObject.remove(key);
+                        basicDBObject.put(queryKey, queryList);
+                        return true;
+                    } else if (valList.size() > 0 && !( (valList.get(0) instanceof BasicDBObject) || (valList.get(0) instanceof BasicDBList) ) && key.equalsIgnoreCase(queryKey)) {
+                        List<Object> queryList = Collections.singletonList(queryVal);
+                        basicDBObject.remove(key);
+                        basicDBObject.put(queryKey, queryList);
+                        return true;
+                    }
+                }
+
+                res = modifyValueInPayload(value, key, queryKey, queryVal);
+                if (res) {
+                    break;
+                }
+            }
+        } else if (obj instanceof BasicDBList) {
+            for(Object elem: (BasicDBList) obj) {
+                res = modifyValueInPayload(elem, parentKey, queryKey, queryVal);
+                if (res) {
+                    break;
+                }
+            }
+        }
+        return res;
+    }
+}
diff --git a/libs/dao/src/main/java/com/akto/dto/testing/HardcodedAuthParam.java b/libs/dao/src/main/java/com/akto/dto/testing/HardcodedAuthParam.java
@@ -1,14 +1,7 @@
 package com.akto.dto.testing;
 
-import com.akto.dto.HttpRequestParams;
 import com.akto.dto.OriginalHttpRequest;
-import com.akto.util.CookieTransformer;
-import com.akto.util.JSONUtils;
-import com.akto.util.JsonStringPayloadModifier;
 import com.akto.util.TokenPayloadModifier;
-import com.mongodb.BasicDBObject;
-
-import java.util.*;
 
 public class HardcodedAuthParam extends AuthParam {
     private Location where;
@@ -40,18 +33,7 @@ public boolean removeAuthTokens(OriginalHttpRequest request) {
 
     @Override
     public boolean authTokenPresent(OriginalHttpRequest request) {
-        if (this.key == null) return false;
-        String k = this.key.toLowerCase().trim();
-
-        if (where.toString().equals(AuthParam.Location.BODY.toString())) {
-            BasicDBObject basicDBObject =  BasicDBObject.parse(request.getBody());
-            BasicDBObject data = JSONUtils.flattenWithDots(basicDBObject);
-            return data.keySet().contains(this.key);
-        } else {
-            Map<String, List<String>> headers = request.getHeaders();
-            List<String> cookieList = headers.getOrDefault("cookie", new ArrayList<>());
-            return headers.containsKey(k) || CookieTransformer.isKeyPresentInCookie(cookieList, k);
-        }
+        return Utils.isRequestKeyPresent(this.key, request, where);
     }
 
     public Location getWhere() {

diff --git a/libs/dao/src/main/java/com/akto/dto/testing/LoginRequestAuthParam.java b/libs/dao/src/main/java/com/akto/dto/testing/LoginRequestAuthParam.java
@@ -1,17 +1,7 @@
 package com.akto.dto.testing;
 
 import com.akto.dto.OriginalHttpRequest;
-import com.akto.util.CookieTransformer;
-import com.akto.util.JSONUtils;
-import com.akto.util.JsonStringPayloadModifier;
 import com.akto.util.TokenPayloadModifier;
-import com.mongodb.BasicDBObject;
-
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
-
 public class LoginRequestAuthParam extends AuthParam {
 
     private Location where;
@@ -43,19 +33,7 @@ public boolean removeAuthTokens(OriginalHttpRequest request) {
 
     @Override
     public boolean authTokenPresent(OriginalHttpRequest request) {
-        if (this.key == null) return false;
-        String k = this.key.toLowerCase().trim();
-
-        if (where.toString().equals(AuthParam.Location.BODY.toString())) {
-            String body = request.getBody();
-            BasicDBObject basicDBObject =  BasicDBObject.parse(request.getBody());
-            BasicDBObject data = JSONUtils.flattenWithDots(basicDBObject);
-            return data.keySet().contains(this.key);
-        } else {
-            Map<String, List<String>> headers = request.getHeaders();
-            List<String> cookieList = headers.getOrDefault("cookie", new ArrayList<>());
-            return headers.containsKey(k) || CookieTransformer.isKeyPresentInCookie(cookieList, k);
-        }
+        return Utils.isRequestKeyPresent(this.key, request, where);
     }
 
     public Location getWhere() {

diff --git a/libs/dao/src/main/java/com/akto/dto/testing/Utils.java b/libs/dao/src/main/java/com/akto/dto/testing/Utils.java
@@ -0,0 +1,36 @@
+package com.akto.dto.testing;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import com.akto.dto.OriginalHttpRequest;
+import com.akto.dto.testing.AuthParam.Location;
+import com.akto.util.CookieTransformer;
+import com.akto.util.JSONUtils;
+import com.mongodb.BasicDBObject;
+
+public class Utils {
+    public static boolean isRequestKeyPresent(String key, OriginalHttpRequest request, Location where){
+        if (key == null) return false;
+        String k = key.toLowerCase().trim();
+        if (where.toString().equals(AuthParam.Location.BODY.toString())) {
+            BasicDBObject basicDBObject =  BasicDBObject.parse(request.getBody());
+            BasicDBObject data = JSONUtils.flattenWithDots(basicDBObject);
+            boolean exists = data.keySet().contains(key);
+            if(!exists){
+                for(String payloadKey: data.keySet()){
+                    if(payloadKey.contains(key)){
+                        exists = true;
+                        break;
+                    }
+                }
+            }
+            return exists;
+        } else {
+            Map<String, List<String>> headers = request.getHeaders();
+            List<String> cookieList = headers.getOrDefault("cookie", new ArrayList<>());
+            return headers.containsKey(k) || CookieTransformer.isKeyPresentInCookie(cookieList, k);
+        }
+    }
+}
diff --git a/libs/dao/src/main/java/com/akto/util/JsonStringPayloadModifier.java b/libs/dao/src/main/java/com/akto/util/JsonStringPayloadModifier.java
@@ -1,35 +1,37 @@
 package com.akto.util;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.mongodb.BasicDBObject;
 
 public class JsonStringPayloadModifier {
     private static final ObjectMapper mapper = new ObjectMapper();
-
     public static String jsonStringPayloadModifier(String data, String path, String newVal) throws Exception {
 
         try {
             JsonNode origRequestNode = mapper.readValue(data, JsonNode.class);
             JsonNode node = origRequestNode;
-            JsonNode parentNode = origRequestNode;
             String []keys = path.split("\\.");
             for (int i=0; i<keys.length - 1; i++) {
                 node = node.get(keys[i]);
-                parentNode = node;
                 if (node == null) {
                     throw new Exception("invalid key specified");
                 }
             }
-            if (node.get(keys[keys.length-1]) == null || !node.get(keys[keys.length-1]).isValueNode()) {
-                throw new Exception("key not found in request payload");
+
+            BasicDBObject payload;
+            try {
+                payload = BasicDBObject.parse(data);
+            } catch (Exception e) {
+                payload = new BasicDBObject();
             }
-
-            if (newVal == null || newVal == "null") {
-                ((ObjectNode) parentNode).remove(keys[keys.length-1]);
-            } else {
-                ((ObjectNode) parentNode).put(keys[keys.length-1], newVal);
+
+            BasicDBObject newValueObj = BasicDBObject.parse("{\"value\":\"" + newVal + "\"}");
+            boolean isModified = com.akto.dto.test_editor.Util.modifyValueInPayload(payload, null, keys[keys.length - 1],newValueObj.get("value"));
+            if(isModified){
+                return payload.toString();
+            }else{
+                return data;
             }
-            return origRequestNode.toString();
 
         } catch (Exception e) {
             throw e;

diff --git a/libs/dao/src/main/java/com/akto/util/TokenPayloadModifier.java b/libs/dao/src/main/java/com/akto/util/TokenPayloadModifier.java
@@ -8,7 +8,6 @@
 import com.akto.dto.OriginalHttpRequest;
 import com.akto.dto.testing.AuthParam;
 
-import okhttp3.Cookie;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 

diff --git a/libs/utils/src/main/java/com/akto/test_editor/Utils.java b/libs/utils/src/main/java/com/akto/test_editor/Utils.java
@@ -1,6 +1,5 @@
 package com.akto.test_editor;
 
-import java.util.*;
 import java.io.IOException;
 import java.net.URI;
 import java.util.ArrayList;
@@ -102,56 +101,7 @@ public static boolean deleteKeyFromPayload(Object obj, String parentKey, String
     }
 
     public static boolean modifyValueInPayload(Object obj, String parentKey, String queryKey, Object queryVal) {
-        boolean res = false;
-        if (obj instanceof BasicDBObject) {
-            BasicDBObject basicDBObject = (BasicDBObject) obj;
-
-            Set<String> keySet = basicDBObject.keySet();
-
-            for(String key: keySet) {
-                if (key == null) {
-                    continue;
-                }
-                Object value = basicDBObject.get(key);
-
-                if (!( (value instanceof BasicDBObject) || (value instanceof BasicDBList) )) {
-                    if (key.equalsIgnoreCase(queryKey)) {
-                        basicDBObject.remove(key);
-                        basicDBObject.put(queryKey, queryVal);
-                        return true;
-                    }
-                }
-
-                if (value instanceof BasicDBList) {
-                    BasicDBList valList = (BasicDBList) value;
-                    if (valList.size() == 0 && key.equalsIgnoreCase(queryKey)) {
-                        List<Object> queryList = Collections.singletonList(queryVal);
-                        basicDBObject.remove(key);
-                        basicDBObject.put(queryKey, queryList);
-                        return true;
-                    } else if (valList.size() > 0 && !( (valList.get(0) instanceof BasicDBObject) || (valList.get(0) instanceof BasicDBList) ) && key.equalsIgnoreCase(queryKey)) {
-                        List<Object> queryList = Collections.singletonList(queryVal);
-                        basicDBObject.remove(key);
-                        basicDBObject.put(queryKey, queryList);
-                        return true;
-                    }
-                }
-
-                res = modifyValueInPayload(value, key, queryKey, queryVal);
-                if (res) {
-                    break;
-                }
-            }
-        } else if (obj instanceof BasicDBList) {
-            for(Object elem: (BasicDBList) obj) {
-                res = modifyValueInPayload(elem, parentKey, queryKey, queryVal);
-                if (res) {
-                    break;
-                }
-            }
-        }
-
-        return res;
+        return com.akto.dto.test_editor.Util.modifyValueInPayload(obj, parentKey, queryKey, queryVal);
     }
 
     public static String applyRegexModifier(String data, String regex, String replaceWith) {