From ca6475b19208c161cda2a20109ca13fe53494616 Mon Sep 17 00:00:00 2001
From: Avneesh Hota <avneesh@akto.io>
Date: Thu, 23 Nov 2023 12:40:58 +0530
Subject: [PATCH] fixed word list append feature

---
 .../execution/VariableResolver.java           | 50 ++++++++++++++++---
 .../execution/TestVariableResolver.java       | 43 ++++++++++++++++
 2 files changed, 86 insertions(+), 7 deletions(-)

diff --git a/apps/testing/src/main/java/com/akto/test_editor/execution/VariableResolver.java b/apps/testing/src/main/java/com/akto/test_editor/execution/VariableResolver.java
index 9fabbdfca9..fc41179d39 100644
--- a/apps/testing/src/main/java/com/akto/test_editor/execution/VariableResolver.java
+++ b/apps/testing/src/main/java/com/akto/test_editor/execution/VariableResolver.java
@@ -220,20 +220,56 @@ public static Boolean isWordListVariable(Object key, Map<String, Object> varMap)
 
         String expression = key.toString();
 
-        expression = expression.substring(2, expression.length());
-        expression = expression.substring(0, expression.length() - 1);
+        Pattern pattern = Pattern.compile("\\$\\{[^}]*\\}");
+        Matcher matcher = pattern.matcher(expression);
+        while (matcher.find()) {
+            try {
+                String match = matcher.group(0);
+                match = match.substring(2, match.length());
+                match = match.substring(0, match.length() - 1);
+
+                Boolean isWordListVar = varMap.containsKey("wordList_" + match);
+                if (isWordListVar) return true;
+            } catch (Exception e) {
+                e.printStackTrace();
+            }
+        }
 
-        Boolean isWordListVar = varMap.containsKey("wordList_" + expression);
-        return isWordListVar;
+        return false;
     }
 
     public static List<String> resolveWordListVar(String key, Map<String, Object> varMap) {
         String expression = key.toString();
 
-        expression = expression.substring(2, expression.length());
-        expression = expression.substring(0, expression.length() - 1);
+        List<String> wordList = new ArrayList<>();
+        String wordListKey = null;
+
+        Pattern pattern = Pattern.compile("\\$\\{[^}]*\\}");
+        Matcher matcher = pattern.matcher(expression);
+        while (matcher.find()) {
+            try {
+                String match = matcher.group(0);
+                String originalKey = match;
+                match = match.substring(2, match.length());
+                match = match.substring(0, match.length() - 1);
+
+                Boolean isWordListVar = varMap.containsKey("wordList_" + match);
+                if (isWordListVar) {
+                    wordList = (List<String>) varMap.get("wordList_" + match);
+                    wordListKey = originalKey;
+                    break;
+                }
+            } catch (Exception e) {
+                e.printStackTrace();
+            }
+        }
+
+        List<String> result = new ArrayList<>();
+        for (String word: wordList) {
+            result.add(expression.replace(wordListKey, word));
+        }
 
-        return (List<String>) varMap.get("wordList_" + expression);
+        return result;
     }
 
     // public Object resolveExpression(Map<String, Object> varMap, String expression) {
diff --git a/apps/testing/src/test/java/com/akto/test_editor/execution/TestVariableResolver.java b/apps/testing/src/test/java/com/akto/test_editor/execution/TestVariableResolver.java
index edc6aac402..1c00d69517 100644
--- a/apps/testing/src/test/java/com/akto/test_editor/execution/TestVariableResolver.java
+++ b/apps/testing/src/test/java/com/akto/test_editor/execution/TestVariableResolver.java
@@ -2,7 +2,9 @@
 
 import org.junit.Test;
 
+import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 import static junit.framework.TestCase.assertEquals;
@@ -44,4 +46,45 @@ public void testResolveExpression() {
         assertEquals("akto", result);
     }
 
+    @Test
+    public void testResolveWordListVar() {
+        Map<String, Object> varMap = new HashMap<>();
+        varMap.put("changed_body_value", "akto");
+        varMap.put("randomVar", "random");
+        varMap.put("wordList_specialCharacters", Arrays.asList(".", "$", "/"));
+        String key = "${changed_body_value}${specialCharacters}${randomVar}";
+
+        List<String> result = VariableResolver.resolveWordListVar(key, varMap);
+        assertEquals(Arrays.asList("${changed_body_value}.${randomVar}", "${changed_body_value}$${randomVar}", "${changed_body_value}/${randomVar}"), result);
+
+
+        varMap = new HashMap<>();
+        varMap.put("changed_body_value", "akto");
+        varMap.put("randomVar", "random");
+        varMap.put("wordList_specialCharacters", Arrays.asList(".", "$", "/"));
+        key = "asdf${specialCharacters}xyz";
+
+        result = VariableResolver.resolveWordListVar(key, varMap);
+        assertEquals(Arrays.asList("asdf.xyz", "asdf$xyz", "asdf/xyz"), result);
+
+        varMap = new HashMap<>();
+        varMap.put("changed_body_value", "akto");
+        varMap.put("randomVar", "random");
+        varMap.put("wordList_specialCharacters", Arrays.asList(".", "$", "/"));
+        key = "${specialCharacters}";
+
+        result = VariableResolver.resolveWordListVar(key, varMap);
+        assertEquals(Arrays.asList(".", "$", "/"), result);
+
+        varMap = new HashMap<>();
+        varMap.put("changed_body_value", "akto");
+        varMap.put("randomVar", "random");
+        varMap.put("wordList_specialCharacters", Arrays.asList(".", "$", "/"));
+        varMap.put("wordList_names", Arrays.asList(".", "$", "/"));
+        key = "${changed_body_value}${specialCharacters}${names}";
+
+        result = VariableResolver.resolveWordListVar(key, varMap);
+        assertEquals(Arrays.asList("${changed_body_value}.${names}", "${changed_body_value}$${names}", "${changed_body_value}/${names}"), result);
+    }
+
 }