Merge pull request #8 from Hannah-PortSwigger/editable-upgrade-request

v1.1.0

build.gradle

@@ -3,7 +3,7 @@ plugins {
 }
 
 group = 'net.portswigger'
-version = '1.0.1'
+version = '1.1.0'
 
 repositories {
     mavenLocal()
@@ -15,12 +15,14 @@ dependencies {
 
     implementation (
             'org.python:jython-standalone:2.7.3',
-            'com.fifesoft:rsyntaxtextarea:3.3.4',
-            'commons-io:commons-io:2.13.0'
+            'com.fifesoft:rsyntaxtextarea:3.3.4'
     )
 }
 
 tasks.register('fatJar', Jar) {
+    manifest {
+        attributes 'implementation-version': "${version}"
+    }
     duplicatesStrategy = DuplicatesStrategy.EXCLUDE
     from { configurations.compileClasspath.collect { it.isDirectory() ? it : zipTree(it) } }
     with jar

src/main/java/attack/AttackHandler.java

@@ -1,6 +1,6 @@
 package attack;
 
-import burp.api.montoya.ui.contextmenu.WebSocketMessage;
+import burp.api.montoya.http.message.requests.HttpRequest;
 import burp.api.montoya.websocket.Direction;
 import burp.api.montoya.websocket.WebSockets;
 import connection.Connection;
@@ -27,7 +27,6 @@ public class AttackHandler
     private final BlockingQueue<WebSocketConnectionMessage> sendMessageQueue;
     private final BlockingQueue<ConnectionMessage> tableBlockingQueue;
     private final WebSocketMessageTableModel webSocketMessageTableModel;
-    private final WebSocketMessage baseWebSocketMessage;
     private final AtomicBoolean isAttackRunning;
     private final PythonInterpreter interpreter;
     private ExecutorService sendMessageExecutorService;
@@ -39,32 +38,30 @@ public AttackHandler(
             BlockingQueue<WebSocketConnectionMessage> sendMessageQueue,
             BlockingQueue<ConnectionMessage> tableBlockingQueue,
             WebSocketMessageTableModel webSocketMessageTableModel,
-            WebSocketMessage baseWebSocketMessage,
             AtomicBoolean isAttackRunning
     )
     {
         this.logger = logger;
         this.sendMessageQueue = sendMessageQueue;
         this.tableBlockingQueue = tableBlockingQueue;
         this.webSocketMessageTableModel = webSocketMessageTableModel;
-        this.baseWebSocketMessage = baseWebSocketMessage;
         this.isAttackRunning = isAttackRunning;
+
         interpreter = new PythonInterpreter();
+
         interpreter.setOut(logger.outputStream());
         interpreter.setErr(logger.errorStream());
 
-        interpreter.set("base_websocket", baseWebSocketMessage);
-
         interpreter.set("websocket_connection", new ConnectionFactory(logger, webSockets, sendMessageQueue, isAttackRunning));
-
         interpreter.set("results_table", new TableBlockingQueueProducer(logger, tableBlockingQueue));
     }
 
-    public void executeJython(String payload, String editorCodeString)
+    public void executeJython(String payload, HttpRequest upgradeRequest, String editorCodeString)
     {
         interpreter.set("payload", payload);
+        interpreter.set("upgrade_request", upgradeRequest);
         interpreter.exec(editorCodeString);
-        interpreter.exec("queue_websockets(base_websocket, payload)");
+        interpreter.exec("queue_websockets(upgrade_request, payload)");
     }
 
     public void executeCallback(WebSocketConnectionMessage webSocketConnectionMessage)
@@ -94,11 +91,6 @@ public WebSocketMessageTableModel getWebSocketMessageTableModel()
         return webSocketMessageTableModel;
     }
 
-    public WebSocketMessage getBaseWebSocketMessage()
-    {
-        return baseWebSocketMessage;
-    }
-
     public AtomicBoolean getIsAttackRunning()
     {
         return isAttackRunning;

src/main/java/burp/WebSocketContextMenuItemsProvider.java

@@ -15,11 +15,11 @@
 
 public class WebSocketContextMenuItemsProvider implements ContextMenuItemsProvider
 {
-    private final List<JFrame> frameList;
     private final Logger logger;
     private final UserInterface userInterface;
     private final Persistence persistence;
     private final WebSockets webSockets;
+    private final List<JFrame> frameList;
 
     public WebSocketContextMenuItemsProvider(
             Logger logger,
@@ -36,6 +36,7 @@ public WebSocketContextMenuItemsProvider(
         this.frameList = frameList;
     }
 
+    @Override
     public List<Component> provideMenuItems(WebSocketContextMenuEvent event)
     {
         JMenuItem sendToContextMenuItem = new JMenuItem("Send to " + WebSocketFuzzer.EXTENSION_NAME);

src/main/java/burp/WebSocketExtensionWebSocketMessageHandler.java

@@ -6,7 +6,6 @@
 import connection.WebSocketConnection;
 import data.WebSocketConnectionMessage;
 import logger.Logger;
-import logger.LoggerLevel;
 
 import java.time.LocalDateTime;
 import java.util.concurrent.BlockingQueue;
@@ -37,13 +36,13 @@ public void textMessageReceived(TextMessage textMessage)
         }
         catch (InterruptedException e)
         {
-            logger.logError(LoggerLevel.ERROR, "Failed to put message on queue.");
+            logger.logError("Failed to put message on queue.");
         }
     }
 
     @Override
     public void binaryMessageReceived(BinaryMessage binaryMessage)
     {
-        logger.logOutput(LoggerLevel.ERROR, "Unhandled binary message received");
+        logger.logError("Unhandled binary message received");
     }
-}
+}

src/main/java/burp/WebSocketFuzzer.java

@@ -8,13 +8,12 @@
 import burp.api.montoya.websocket.WebSockets;
 import logger.Logger;
 import logger.LoggerLevel;
+import utils.Utilities;
 
 import javax.swing.*;
 import java.util.ArrayList;
 import java.util.List;
 
-import static utils.Utilities.closeAllFrames;
-
 public class WebSocketFuzzer implements BurpExtension
 {
     public static final String EXTENSION_NAME = "WebSocket Turbo Intruder";
@@ -26,56 +25,22 @@ public void initialize(MontoyaApi api)
         Extension extension = api.extension();
         Persistence persistence = api.persistence();
         UserInterface userInterface = api.userInterface();
-        Logger logger = new Logger(api.logging());
         WebSockets websockets = api.websockets();
 
-        List<JFrame> frameList = new ArrayList<>();
+        Logger logger = new Logger(api.logging());
 
-        extension.setName(EXTENSION_NAME);
+        Utilities.initializeDefaultDirectory(logger, persistence);
 
-        initializeDefaultDirectory(logger, persistence);
+        List<JFrame> frameList = new ArrayList<>();
 
-        JMenu menu = generateMenu(logger, persistence, frameList);
+        JMenu menu = Utilities.generateMenu(logger, persistence, frameList);
         userInterface.menuBar().registerMenu(menu);
 
         userInterface.registerContextMenuItemsProvider(new WebSocketContextMenuItemsProvider(logger, userInterface, persistence, websockets, frameList));
-
         extension.registerUnloadingHandler(new WebSocketExtensionUnloadingHandler(frameList));
 
-        logger.logOutput(LoggerLevel.DEFAULT, EXTENSION_NAME + " - Loaded");
-    }
-
-    private void initializeDefaultDirectory(Logger logger, Persistence persistence)
-    {
-        if (persistence.preferences().getString("websocketsScriptsPath") == null)
-        {
-            persistence.preferences().setString("websocketsScriptsPath", DEFAULT_SCRIPT_DIRECTORY);
-            logger.logOutput(LoggerLevel.DEBUG, "Default script directory initialized.");
-        }
-    }
-
-    private JMenu generateMenu(Logger logger, Persistence persistence, List<JFrame> frameList)
-    {
-        JMenuItem resetDefaultScriptsMenuItem = new JMenuItem("Reset scripts directory to default.");
-        resetDefaultScriptsMenuItem.addActionListener(l -> {
-            persistence.preferences().setString("websocketsScriptsPath", DEFAULT_SCRIPT_DIRECTORY);
-            logger.logOutput(LoggerLevel.DEBUG, "Scripts directory reset to " + DEFAULT_SCRIPT_DIRECTORY);
-        });
-
-        JMenuItem closeAllFramesMenuItem = new JMenuItem("Close all " + EXTENSION_NAME + " windows.");
-        closeAllFramesMenuItem.addActionListener(l -> {
-            closeAllFrames(frameList);
-            logger.logOutput(LoggerLevel.DEBUG, "All " + EXTENSION_NAME + " windows closed.");
-        });
-
-        JCheckBoxMenuItem loggingLevelDebug = new JCheckBoxMenuItem("Debug mode", logger.isDebugLogLevel());
-        loggingLevelDebug.addActionListener(l -> logger.setDebugLogLevel(loggingLevelDebug.getState()));
-
-        JMenu menu = new JMenu(EXTENSION_NAME);
-        menu.add(resetDefaultScriptsMenuItem);
-        menu.add(closeAllFramesMenuItem);
-        menu.add(loggingLevelDebug);
-
-        return menu;
+        extension.setName(EXTENSION_NAME);
+        String extensionVersion = WebSocketFuzzer.class.getPackage().getImplementationVersion();
+        logger.logOutput(LoggerLevel.DEFAULT, EXTENSION_NAME + " v" + extensionVersion);
     }
 }

src/main/java/connection/Connection.java

@@ -1,8 +1,12 @@
 package connection;
 
+import burp.api.montoya.http.message.requests.HttpRequest;
+
 public interface Connection
 {
     void queue(String payload);
 
     void queue(String payload, String comment);
+
+    HttpRequest upgradeRequest();
 }

src/main/java/connection/ConnectionFactory.java

@@ -1,6 +1,6 @@
 package connection;
 
-import burp.api.montoya.ui.contextmenu.WebSocketMessage;
+import burp.api.montoya.http.message.requests.HttpRequest;
 import burp.api.montoya.websocket.WebSockets;
 import data.WebSocketConnectionMessage;
 import logger.Logger;
@@ -28,8 +28,8 @@ public ConnectionFactory(
         this.isAttackRunning = isAttackRunning;
     }
 
-    public Connection create(WebSocketMessage baseWebSocketMessage)
+    public Connection create(HttpRequest upgradeRequest)
     {
-        return new WebSocketConnection(logger, webSockets, sendMessageQueue, baseWebSocketMessage, isAttackRunning);
+        return new WebSocketConnection(logger, webSockets, sendMessageQueue, upgradeRequest, isAttackRunning);
     }
 }

src/main/java/connection/WebSocketConnection.java

@@ -1,7 +1,7 @@
 package connection;
 
 import burp.WebSocketExtensionWebSocketMessageHandler;
-import burp.api.montoya.ui.contextmenu.WebSocketMessage;
+import burp.api.montoya.http.message.requests.HttpRequest;
 import burp.api.montoya.websocket.Direction;
 import burp.api.montoya.websocket.WebSockets;
 import burp.api.montoya.websocket.extension.ExtensionWebSocket;
@@ -19,23 +19,25 @@ public class WebSocketConnection implements Connection
     private final Logger logger;
     private final WebSockets webSockets;
     private final BlockingQueue<WebSocketConnectionMessage> sendMessageQueue;
+    private final HttpRequest upgradeRequest;
     private final AtomicBoolean isAttackRunning;
     private final ExtensionWebSocket extensionWebSocket;
 
     WebSocketConnection(
             Logger logger,
             WebSockets webSockets,
             BlockingQueue<WebSocketConnectionMessage> sendMessageQueue,
-            WebSocketMessage baseWebSocketMessage,
+            HttpRequest upgradeRequest,
             AtomicBoolean isAttackRunning
     )
     {
         this.logger = logger;
         this.webSockets = webSockets;
         this.sendMessageQueue = sendMessageQueue;
+        this.upgradeRequest = upgradeRequest;
         this.isAttackRunning = isAttackRunning;
 
-        extensionWebSocket = createExtensionWebSocket(baseWebSocketMessage);
+        extensionWebSocket = createExtensionWebSocket(upgradeRequest);
     }
 
     @Override
@@ -48,7 +50,7 @@ public void queue(String payload)
             }
             catch (InterruptedException e)
             {
-                logger.logError(LoggerLevel.ERROR, "Failed to put message on sendMessageQueue");
+                logger.logError("Failed to put message on sendMessageQueue");
             }
         }
     }
@@ -57,27 +59,35 @@ public void queue(String payload)
     public void queue(String payload, String comment)
     {
         if (isAttackRunning.get())
-        {try
+        {
+            try
             {
                 sendMessageQueue.put(new WebSocketConnectionMessage(payload, Direction.CLIENT_TO_SERVER, LocalDateTime.now(), comment, this));
             }
             catch (InterruptedException e)
             {
-                logger.logError(LoggerLevel.ERROR, "Failed to put message on sendMessageQueue");
+                logger.logError("Failed to put message on sendMessageQueue");
             }
         }
     }
 
+    @Override
+    public HttpRequest upgradeRequest()
+    {
+        return upgradeRequest;
+    }
+
     public void sendMessage(String payload)
     {
         extensionWebSocket.sendTextMessage(payload);
     }
 
-    private ExtensionWebSocket createExtensionWebSocket(WebSocketMessage baseWebSocketMessage)
+    private ExtensionWebSocket createExtensionWebSocket(HttpRequest upgradeRequest)
     {
         ExtensionWebSocket extensionWebSocket;
 
-        ExtensionWebSocketCreation extensionWebSocketCreation = webSockets.createWebSocket(baseWebSocketMessage.upgradeRequest());
+        ExtensionWebSocketCreation extensionWebSocketCreation = webSockets.createWebSocket(upgradeRequest);
+        logger.logOutput(LoggerLevel.DEBUG, "WebSocketConnection Upgrade request: " + upgradeRequest.toString());
 
         if (extensionWebSocketCreation.webSocket().isPresent())
         {
@@ -87,7 +97,7 @@ private ExtensionWebSocket createExtensionWebSocket(WebSocketMessage baseWebSock
         }
         else
         {
-            logger.logError(LoggerLevel.DEFAULT, "Failed to create websocket connection");
+            logger.logError("Failed to create websocket connection");
             extensionWebSocket = null;
         }
 

src/main/java/logger/Logger.java

@@ -5,20 +5,16 @@
 import java.io.OutputStream;
 
 import static logger.LoggerLevel.DEBUG;
-import static logger.LoggerLevel.ERROR;
 
 public class Logger
 {
     private final Logging logging;
     private boolean debugLogLevel;
-    private final boolean errorLogLevel;
 
     public Logger(Logging logging)
     {
         this.logging = logging;
-
         debugLogLevel  = false;
-        errorLogLevel = true;
     }
 
     public boolean isDebugLogLevel()
@@ -39,12 +35,9 @@ public void logOutput(LoggerLevel loggerLevel, String output)
         }
     }
 
-    public void logError(LoggerLevel loggerLevel, String output)
+    public void logError(String output)
     {
-        if (isValid(loggerLevel))
-        {
-            logging.logToError(output);
-        }
+        logging.logToError(output);
     }
 
     public OutputStream outputStream()
@@ -59,10 +52,6 @@ public OutputStream errorStream()
 
     private boolean isValid(LoggerLevel outputLogLevel)
     {
-        if (outputLogLevel.equals(DEBUG) && !debugLogLevel)
-        {
-            return false;
-        }
-        else return !outputLogLevel.equals(ERROR) || errorLogLevel;
+        return !outputLogLevel.equals(DEBUG) || debugLogLevel;
     }
 }