modified: src/main/java/org.focalpoint.isns.burp.srichecks/ScriptFin…

…der.java

src/main/java/org.focalpoint.isns.burp.srichecks/ScriptFinder.java

@@ -29,6 +29,8 @@
 import org.openqa.selenium.StaleElementReferenceException;
 import org.openqa.selenium.TimeoutException;
 
+import org.openqa.selenium.Cookie;
+
 import org.openqa.selenium.chrome.ChromeDriverService;
 import org.openqa.selenium.remote.RemoteWebDriver;
 import java.io.File;
@@ -52,13 +54,15 @@
 import java.net.URI;
 import java.net.MalformedURLException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 
 public class ScriptFinder{
     private IBurpExtenderCallbacks myCallbacks;
     private Integer PAGE_WAIT_TIMEOUT = 10;
     private String url="NONE";
     private String html="NONE";
+    private List<String> requestHeaders = new ArrayList<>();
     private List<String> domScripts = new ArrayList<>();
     private List<String> htmlScripts = new ArrayList<>();
     // Something to store a parsed URL
@@ -138,6 +142,14 @@ public Integer getTimeout(){
         return PAGE_WAIT_TIMEOUT;
     }
 
+    /**
+     * Set the request headers
+     * @param headers - a list of request headers
+     */
+    public void setRequestHeaders(List<String> headers){
+        Collections.copy(requestHeaders, headers);
+    }
+
 
     /**
      * There is no reason that this should ever be called within burp. It is just here for tests.
@@ -173,6 +185,7 @@ public void startDriver(){
             HashMap<String, Object> prefs = new HashMap<String, Object>(); 
             prefs.put("profile.managed_default_content_settings.images", 2);
             options.setExperimentalOption("prefs", prefs); 
+
             driver = new RemoteWebDriver(serviceManager.getService().getUrl(), options);
             driver.manage().timeouts().implicitlyWait(PAGE_WAIT_TIMEOUT, TimeUnit.SECONDS); // Wait for the page to be completely loaded. Or reasonably loaded.
         }
@@ -187,7 +200,24 @@ public void startDriver(){
      */
     public void checkForDomScripts(){
         startDriver();
-        try{
+
+        // Set the driver's cookies based on the headers, if there are any
+        if (requestHeaders != null){
+            for (String header: requestHeaders){
+                if (header.startsWith("Cookie: ")){
+                    // This is a cookie, split it up
+                    String cookieString = header.substring(8,header.length());
+                    for (String kvPair : cookieString.split(";")){
+                        String key = kvPair.split("=")[0];
+                        String value = kvPair.split("=")[1];
+                        Cookie cookieObj = new Cookie(key, value);
+                        driver.manage().setCookie(cookieObj);
+                    }
+                }
+            }
+        }
+
+        try {
             driver.get(url);
         }
         catch (TimeoutException e){