Add sequence args

.gitignore

@@ -8,4 +8,5 @@ releases/
 *.iml
 gradle-app.setting
 .gradletasknamecache
-.git
+.git
+bin

src/main/java/com/coreyd97/stepper/MessageProcessor.java

@@ -1,11 +1,15 @@
 package com.coreyd97.stepper;
 
 import burp.*;
+
 import com.coreyd97.BurpExtenderUtilities.Preferences;
 import com.coreyd97.stepper.sequence.StepSequence;
 import com.coreyd97.stepper.sequencemanager.SequenceManager;
 import com.coreyd97.stepper.util.ReplacingInputStream;
+import com.coreyd97.stepper.variable.PreExecutionStepVariable;
 import com.coreyd97.stepper.variable.StepVariable;
+import com.google.gson.Gson;
+import com.google.gson.reflect.TypeToken;
 
 import javax.swing.*;
 import java.io.ByteArrayInputStream;
@@ -22,13 +26,20 @@ public class MessageProcessor implements IHttpListener {
     private final Preferences preferences;
     public static final String EXECUTE_BEFORE_HEADER = "X-Stepper-Execute-Before";
     public static final String EXECUTE_AFTER_HEADER = "X-Stepper-Execute-After";
+    public static final String EXECUTE_VAR_HEADER = "X-Stepper-Argument";
     public static final String EXECUTE_BEFORE_REGEX = EXECUTE_BEFORE_HEADER + ":(.*)";
     public static final String EXECUTE_AFTER_REGEX = EXECUTE_AFTER_HEADER+":(.*)";
+    public static final String EXECUTE_VAR_REGEX = EXECUTE_VAR_HEADER + ":(.*)";
     public static final String EXECUTE_AFTER_COMMENT_DELIMITER = "#%~%#";
     public static final Pattern EXECUTE_BEFORE_HEADER_PATTERN = Pattern.compile("^" + EXECUTE_BEFORE_REGEX + "$", Pattern.CASE_INSENSITIVE);
     public static final Pattern EXECUTE_AFTER_HEADER_PATTERN = Pattern.compile("^" + EXECUTE_AFTER_REGEX + "$", Pattern.CASE_INSENSITIVE);
+    public static final Pattern EXECUTE_VAR_HEADER_PATTERN = Pattern.compile("^" + EXECUTE_VAR_REGEX + "$", Pattern.CASE_INSENSITIVE);
     public static final String STEPPER_IGNORE_HEADER = "X-Stepper-Ignore";
     public static final Pattern STEPPER_IGNORE_PATTERN = Pattern.compile("^"+STEPPER_IGNORE_HEADER, Pattern.CASE_INSENSITIVE);
+    public static final Pattern STEPPER_SEQUENCE_NAME_PATTERN = Pattern.compile("^([^:]+)(?::?)", Pattern.CASE_INSENSITIVE);
+    public static final Pattern VARIABLE_LIST_PATTERN = Pattern.compile("[^:]+(:\\s*(?<variables>.+))?", Pattern.CASE_INSENSITIVE);
+    public static final Pattern VARIABLE_PATTERN = Pattern.compile("(?<key>[^=]+)=(?<value>[^;]+);?", Pattern.CASE_INSENSITIVE);
+    public static final Pattern SINGLE_VARIABLE_PATTERN = Pattern.compile("(?<key>[^=]+)=(?<value>.*)", Pattern.CASE_INSENSITIVE);
 
     public MessageProcessor(SequenceManager sequenceManager, Preferences preferences){
         this.sequenceManager = sequenceManager;
@@ -60,32 +71,58 @@ public void processHttpMessage(int toolFlag, boolean messageIsRequest, IHttpRequ
 
             if(messageIsRequest){
                 byte[] request = messageInfo.getRequest();
-                System.out.println("Request: " + messageInfo.getRequest());
-                List<StepSequence> preExecSequences = extractExecSequencesFromRequest(requestInfo, EXECUTE_BEFORE_HEADER_PATTERN);
+
+                // Extract variable headers from request
+                Map<String, String> sequenceArguments = extractSequenceArgumentsFromRequest(requestInfo, EXECUTE_VAR_HEADER_PATTERN);
+
+                List<RequestSequenceInformation> preExecSequences = extractExecSequencesFromRequest(requestInfo, EXECUTE_BEFORE_HEADER_PATTERN);
                 if(preExecSequences.size() > 0){
                     //Remove the headers from the request
                     request = removeHeaderMatchingPattern(request, EXECUTE_BEFORE_HEADER_PATTERN);
 
                     //Execute the sequences
-                    for (StepSequence sequence : preExecSequences) {
-                        sequence.executeBlocking();
+                    for (RequestSequenceInformation requestSequenceInformation : preExecSequences) {
+                        StepSequence sequence = requestSequenceInformation.sequence;
+                        Map<String, String> variables = requestSequenceInformation.variables;
+
+                        // Merge arguments from variable headers into arguments for this sequence
+                        // Prioritize arguments specifically set for this sequence
+                        for (Map.Entry<String, String> entry : sequenceArguments.entrySet()) {
+                            variables.putIfAbsent(entry.getKey(), entry.getValue());
+                        }
+
+                        sequence.executeBlocking(variables);
                     }
                 }
 
-                List<StepSequence> postExecSequences = extractExecSequencesFromRequest(requestInfo, EXECUTE_AFTER_HEADER_PATTERN);
+                List<RequestSequenceInformation> postExecSequences = extractExecSequencesFromRequest(requestInfo, EXECUTE_AFTER_HEADER_PATTERN);
 
                 if(postExecSequences.size() > 0){
                     //Remove the headers from the request
                     request = removeHeaderMatchingPattern(request, EXECUTE_AFTER_HEADER_PATTERN);
 
+                    HashMap<String, Map<String, String>> allSequences = new HashMap<>();
+
                     //Joining the sequences as string to set them as a comment to catch them in the response
                     String postExecSequencesJoined = EXECUTE_AFTER_HEADER + ":";
-                    for (StepSequence sequence : postExecSequences) {
-                        postExecSequencesJoined += sequence.getTitle() + EXECUTE_AFTER_COMMENT_DELIMITER;
+                    for (RequestSequenceInformation requestSequenceInformation : postExecSequences) {
+                        StepSequence sequence = requestSequenceInformation.sequence;
+                        Map<String, String> variables = requestSequenceInformation.variables;
+
+                        // Merge arguments from variable headers into arguments for this sequence
+                        // Prioritize arguments specifically set for this sequence
+                        for (Map.Entry<String, String> entry : sequenceArguments.entrySet()) {
+                            variables.putIfAbsent(entry.getKey(), entry.getValue());
+                        }
+
+                        allSequences.put(sequence.getTitle(), variables);
                     }
 
+                    Gson gson = new Gson();
+                    String serializedSequences = gson.toJson(allSequences);
+
                     if(!postExecSequencesJoined.isEmpty()){
-                        messageInfo.setComment(messageInfo.getComment() + postExecSequencesJoined);
+                        messageInfo.setComment(messageInfo.getComment() + EXECUTE_AFTER_HEADER + ":" + serializedSequences);
                     }
                 }
 
@@ -116,11 +153,13 @@ public void processHttpMessage(int toolFlag, boolean messageIsRequest, IHttpRequ
                 messageInfo.setRequest(request);
             }else{
                 // this is a response so we check for the comments
-                List<StepSequence> postExecSequences = extractExecSequencesFromComment(messageInfo.getComment(), EXECUTE_AFTER_HEADER_PATTERN);
+                List<RequestSequenceInformation> postExecSequences = extractExecSequencesFromComment(messageInfo.getComment(), EXECUTE_AFTER_HEADER_PATTERN);
                 if(postExecSequences.size() > 0){
                     //Execute the sequences
-                    for (StepSequence sequence : postExecSequences) {
-                        sequence.executeBlocking();
+                    for (RequestSequenceInformation requestSequenceInformation : postExecSequences) {
+                        StepSequence sequence = requestSequenceInformation.sequence;
+                        Map<String, String> variables = requestSequenceInformation.variables;
+                        sequence.executeBlocking(variables);
                     }
                     // remove the added comment from the request
                     messageInfo.setComment(messageInfo.getComment().replaceAll(EXECUTE_AFTER_REGEX+EXECUTE_AFTER_COMMENT_DELIMITER,""));
@@ -225,29 +264,135 @@ public static byte[] updateContentLength(byte[] request){
         return Stepper.callbacks.getHelpers().buildHttpMessage(newRequestHeaders, newBody);
     }
 
+    /**
+     * Extract all defined variables from the request sequence information string.
+     * 
+     * If a variable has more than one definition, the later one is used. 
+     * It expects the string to have a format that matches the definition in RequesteSequenceInformation.
+     * 
+     * For example, the string: `test: firstVariable=old; secondVariable=someValue; firstVariable=new`
+     * would be converted to a map with the following content:
+     * - firstVariable: new
+     * - secondVariable: someValue
+     * 
+     * @see RequestSequenceInformation
+     * @param sequenceInfoString The string that contains sequence information
+     * @return All defined variables
+     */
+    public Map<String, String> extractVariablesFromSequenceInfoString(String sequenceInfoString) {
+        HashMap<String, String> map = new HashMap<>();
+
+        Matcher variableListMatcher = VARIABLE_LIST_PATTERN.matcher(sequenceInfoString);
+
+        if (variableListMatcher.find() && variableListMatcher.group("variables") != null) {
+            String variableList = variableListMatcher.group("variables");
+            Matcher variablesMatcher = VARIABLE_PATTERN.matcher(variableList);
+
+            while (variablesMatcher.find()) {
+                String variableKey = variablesMatcher.group("key");
+                String variableValue = variablesMatcher.group("value");
+
+                map.put(variableKey, variableValue);
+            }
+        }
+
+        return map;
+    }
+
+    /**
+     * Extract a sequence's name from the request info string.
+     * @param sequenceInfoString String that contains the name and variables of a step sequence
+     * @return The sequence name
+     */
+    private Optional<String> extractSequenceNameFromSequenceInfoString(String sequenceInfoString) {
+        Matcher nameMatcher = STEPPER_SEQUENCE_NAME_PATTERN.matcher(sequenceInfoString);
+        if (!nameMatcher.find()) {
+            return Optional.empty();
+        }
+
+        return Optional.of(nameMatcher.group(1).trim());
+    }
+
+    /**
+     * Extract arguments from variable headers from the given request.
+     * @param requestInfo The request
+     * @param pattern Pattern used to identify variable header
+     * @return Map of all sequence arguments
+     */
+    private Map<String, String> extractSequenceArgumentsFromRequest(IRequestInfo requestInfo, Pattern pattern) {
+        //Check if headers ask us to execute a request before the request.
+        List<String> requestHeaders = requestInfo.getHeaders();
+        Map<String, String> arguments = new HashMap<>();
+
+        for (Iterator<String> iterator = requestHeaders.iterator(); iterator.hasNext(); ) {
+            String header = iterator.next();
+
+            Matcher m = pattern.matcher(header);
+            if (!m.matches()) {
+                continue;
+            }
+
+            String variableInfo = m.group(1).trim();
+
+
+            Matcher argumentMatcher = SINGLE_VARIABLE_PATTERN.matcher(variableInfo);
+            if (!argumentMatcher.matches()) {
+                continue;
+            }
+
+            String variableKey = argumentMatcher.group("key");
+            String variableValue = argumentMatcher.group("value");
+
+            arguments.put(variableKey, variableValue);
+        }
+
+        return arguments;
+    }
+
     /**
      * Locates the X-Stepper-Execute-Before or X-Stepper-Execute-After headers and returns the matching sequences.
      * @param requestInfo
      * @param pattern
      * @return Optional value of step sequence to execute before or after the request.
      */
-    public List<StepSequence> extractExecSequencesFromRequest(IRequestInfo requestInfo, Pattern pattern){
+    public List<RequestSequenceInformation> extractExecSequencesFromRequest(IRequestInfo requestInfo, Pattern pattern){
         //Check if headers ask us to execute a request before the request.
         List<String> requestHeaders = requestInfo.getHeaders();
-        ArrayList<StepSequence> execSequences = new ArrayList<>();
+        ArrayList<RequestSequenceInformation> execSequences = new ArrayList<>();
 
         for (Iterator<String> iterator = requestHeaders.iterator(); iterator.hasNext(); ) {
             String header = iterator.next();
+
             Matcher m = pattern.matcher(header);
-            if (m.matches()) {
-                Optional<StepSequence> execSequence = sequenceManager.getSequences().stream()
-                        .filter(sequence -> sequence.getTitle().equalsIgnoreCase(m.group(1).trim()))
-                        .findFirst();
+            if (!m.matches()) {
+                continue;
+            }
 
-                if(execSequence.isPresent())
-                    execSequences.add(execSequence.get());
-                else
-                    JOptionPane.showMessageDialog(Stepper.getUI().getUiComponent(), "Could not find execution sequence named: \"" + m.group(1).trim() + "\".");
+            ArrayList<StepSequence> currentSequences = new ArrayList<>();
+            String stepperHeader = m.group(1).trim();
+
+            Optional<String> optionalName = this.extractSequenceNameFromSequenceInfoString(stepperHeader);
+
+            if (optionalName.isEmpty()) {
+                continue;
+            }
+
+            String sequenceName = optionalName.get();
+
+            Optional<StepSequence> execSequence = sequenceManager.getSequences().stream()
+                    .filter(sequence -> sequence.getTitle().equalsIgnoreCase(sequenceName))
+                    .findFirst();
+
+            if(execSequence.isPresent())
+                currentSequences.add(execSequence.get());
+            else
+                JOptionPane.showMessageDialog(Stepper.getUI().getUiComponent(), "Could not find execution sequence named: \"" + sequenceName + "\".");
+
+            // Extract variables and store them in the sequences
+            Map<String, String> arguments = extractVariablesFromSequenceInfoString(stepperHeader);
+
+            for (StepSequence sequence : currentSequences) {
+                execSequences.add(new RequestSequenceInformation(sequence, arguments));
             }
         }
         return execSequences;
@@ -259,22 +404,28 @@ public List<StepSequence> extractExecSequencesFromRequest(IRequestInfo requestIn
      * @param pattern
      * @return Optional value of step sequence to execute after the request.
      */
-    public List<StepSequence> extractExecSequencesFromComment(String comment, Pattern pattern){
-        ArrayList<StepSequence> execSequences = new ArrayList<>();
+    public List<RequestSequenceInformation> extractExecSequencesFromComment(String comment, Pattern pattern){
+        ArrayList<RequestSequenceInformation> execSequences = new ArrayList<>();
         Matcher m = pattern.matcher(comment);
         if (m.find()) {
-            String[] allSequences = m.group(1).split(EXECUTE_AFTER_COMMENT_DELIMITER);
-            for(String sequenceName : allSequences){
-                if(sequenceName != null && !sequenceName.isEmpty()){
-                    Optional<StepSequence> execSequence = sequenceManager.getSequences().stream()
-                            .filter(sequence -> sequence.getTitle().equalsIgnoreCase(sequenceName.trim()))
-                            .findFirst();
-
-                    if(execSequence.isPresent())
-                        execSequences.add(execSequence.get());
-                    else
-                        JOptionPane.showMessageDialog(Stepper.getUI().getUiComponent(), "Could not find execution sequence named: \"" + m.group(1).trim() + "\".");
-                }
+            String allSequencesSerialized = m.group(1);
+            Gson gson = new Gson();
+            TypeToken<HashMap<String, HashMap<String, String>>> mapType = new TypeToken<HashMap<String, HashMap<String, String>>>(){};
+            HashMap<String, HashMap<String, String>> allSequences = gson.fromJson(allSequencesSerialized, mapType.getType());
+
+            for(Map.Entry<String, HashMap<String, String>> entry : allSequences.entrySet()){
+                String sequenceName = entry.getKey();
+
+                Optional<StepSequence> execSequence = sequenceManager.getSequences().stream()
+                        .filter(sequence -> sequence.getTitle().equalsIgnoreCase(sequenceName))
+                        .findFirst();
+
+                Map<String, String> variables = entry.getValue();
+
+                if(execSequence.isPresent())
+                    execSequences.add(new RequestSequenceInformation(execSequence.get(), variables));
+                else
+                    JOptionPane.showMessageDialog(Stepper.getUI().getUiComponent(), "Could not find execution sequence named: \"" + sequenceName + "\".");
             }
         }
 

src/main/java/com/coreyd97/stepper/RequestSequenceInformation.java

@@ -0,0 +1,43 @@
+package com.coreyd97.stepper;
+
+import java.util.Map;
+
+import com.coreyd97.stepper.sequence.StepSequence;
+
+/**
+ * Holds information about a step sequence as defined in a request with the stepper headers.
+ * Includes information about all defined variables.
+ * 
+ * This information is stored either in the request's X-Stepper-Execute-Before and X-Steper-Execute-After header
+ * as well in the user comments in the following format:
+ * <sequence_name>: (<key>=<value>;)*
+ * 
+ * Here some examples to demonstrate how the string representation of this looks like:
+ * <pre>
+ * - `simpleExample: first=value`
+ *  - Name: simpleExample
+ *  - Variables:
+ *      first: value
+ * - `advancedExample: firstKey=firstValue;secondKey=secondValue; firstKey=overwrittenValue;`
+ *  - Name: advancedExample
+ *  - Variables:
+ *      firstKey: overwrittenValue
+ *      secondKey: secondValue
+ * </pre>
+ */
+class RequestSequenceInformation {
+    /**
+     * Defines step sequence.
+     */
+    public StepSequence sequence;
+
+    /**
+     * All variables for the sequence as defined in the stepper header.
+     */
+    public Map<String, String> variables;
+
+    public RequestSequenceInformation(StepSequence sequence, Map<String, String> variables) {
+        this.sequence = sequence;
+        this.variables = variables;
+    }
+}