Merge pull request #1995 from akto-api-security/regex_replace_wordlis…

…t_fix add dynamic integer wordlist support
akto-api-security · Jan 17, 2025 · a785baf · a785baf
2 parents ed672ab + 63d15c7
commit a785baf
Showing 1 changed file with 42 additions and 0 deletions.
diff --git a/libs/utils/src/main/java/com/akto/test_editor/execution/VariableResolver.java b/libs/utils/src/main/java/com/akto/test_editor/execution/VariableResolver.java
@@ -6,6 +6,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Random;
 import java.util.Set;
 import java.util.UUID;
 import java.util.regex.Matcher;
@@ -473,6 +474,18 @@ public static Map<String, List<String>> resolveWordList(Map<String, List<String>
                 }
             }
 
+            if (m.containsKey("count")) {
+                int count;
+                if (m.get("count") instanceof String) {
+                    count = Integer.valueOf(m.get("count"));
+                } else {
+                    count = Integer.valueOf((String.valueOf(m.get("count"))));
+                }
+                List<String> vals = generateUniqueTenDigitNumbers(count);
+                wordListsMap.put(k, vals);
+                continue;
+            }
+
             keyObj = m.get("key");
             location = m.get("location");
             if (keyObj instanceof Map) {
@@ -524,6 +537,24 @@ public static Map<String, List<String>> resolveWordList(Map<String, List<String>
 
     }
 
+    public static List<String> generateUniqueTenDigitNumbers(int count) {
+        if (count > 1_000_000_000) {
+            throw new IllegalArgumentException("Cannot generate more than 1 billion unique 10-digit numbers.");
+        }
+
+        List<String> uniqueNumbers = new ArrayList<>();
+        Random random = new Random();
+
+        while (uniqueNumbers.size() < count) {
+            // Generate a 10-digit number starting with 1
+            int number = 1_000_000_000 + random.nextInt(1_000_000_000); // Ensures the number starts with 1
+            uniqueNumbers.add(String.valueOf(number));
+        }
+
+        return uniqueNumbers;
+    }
+
+
     public static Map<String, Object> resolveDynamicWordList(Map<String, Object> varMap, ApiInfo.ApiInfoKey apiInfoKey, Map<ApiInfo.ApiInfoKey, List<String>> newSampleDataMap) {
 
         Map<String, Object> updatedVarMap = new HashMap<>();
@@ -623,6 +654,17 @@ public static void resolveWordList(Map<String, Object> varMap, Map<ApiInfoKey, L
                     continue;
                 }
             }
+            if (m.containsKey("count")) {
+                int count;
+                if (m.get("count") instanceof String) {
+                    count = Integer.valueOf(m.get("count"));
+                } else {
+                    count = Integer.valueOf((String.valueOf(m.get("count"))));
+                }
+                List<String> vals = generateUniqueTenDigitNumbers(count);
+                varMap.put(k, vals);
+                continue;
+            }
             keyObj = m.get("key");
             if (keyObj instanceof Map) {
                 Map<String, String> kMap = (Map) keyObj;