From 5a02b6acf08424bc5391fa72dbc8013d9fc88b78 Mon Sep 17 00:00:00 2001
From: Hannah Law <hannah.law@portswigger.net>
Date: Fri, 10 Nov 2023 14:21:45 +0000
Subject: [PATCH] Add documentation. Add ability to setComment on queue
 messages.

---
 README.md                                     | 41 +++++++++++++++++++
 src/main/java/connection/Connection.java      |  2 +
 .../java/connection/WebSocketConnection.java  | 15 +++++++
 .../examples/BasicExampleWithComments.py      | 12 ++++++
 4 files changed, 70 insertions(+)
 create mode 100644 README.md
 create mode 100644 src/main/resources/examples/BasicExampleWithComments.py

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..ab62d00
--- /dev/null
+++ b/README.md
@@ -0,0 +1,41 @@
+# WebSocket Turbo Intruder
+Extension to fuzz WebSocket messages using custom Python code
+
+## Usage
+1. Right-click on a WebSockets message and go to `Extensions > WebSocket Turbo Intruder > Send to WebSocket Turbo Intruder`
+2. Select a template from the drop-down list
+3. Adjust Python code to suit your use case
+4. Start attack
+
+Note: This will use a new WebSocket connection to send messages down.
+
+## Documentation
+
+### `queue_websockets(base_websocket, payload)`
+`websocket_connection`: This object has one available method - `create()`. Use this to create a WebSocket connection.
+The `create` method takes `base_websocket` as an argument.
+
+Once you've created your WebSocket connection, you can queue messages to send down this connection.
+
+Use the `queue()` method on this object.
+- `queue(String payload)`: Send payload with no comment set
+- `queue(String payload, String comment)`: Send payload with custom comment
+
+`payload` is the contents of the WebSocket message editor in the top half of your screen. You can manually change this, or you can manipulate the String contents in your Python code.
+
+### `handle_outgoing_message(websocket_message)`
+Use this method to conditionally add outgoing messages to the results table.
+
+### `handle_incoming_message(websocket_message)`
+Use this method to conditionally add incoming messages to the results table.
+
+
+### `websocket_message`
+Methods:
+- `getPayload()`: Retrieve the String payload that was sent/received
+- `getDirection()`: Retrieves a `burp.api.montoya.websocket.Direction`
+- `getLength()`: Retrieves the length of the message
+- `getDateTime()`: Retrieves the `java.time.LocalDateTime` that was set on the object
+- `getComment()`: Retrieves the comment that was set on the message
+- `setComment(String comment)`: Allows you to set a comment on the object
+- `getConnection()`: Retrieves the Connection so that you can `queue()` additional messages
\ No newline at end of file
diff --git a/src/main/java/connection/Connection.java b/src/main/java/connection/Connection.java
index 52f309c..cf7482b 100644
--- a/src/main/java/connection/Connection.java
+++ b/src/main/java/connection/Connection.java
@@ -3,4 +3,6 @@
 public interface Connection
 {
     void queue(String payload);
+
+    void queue(String payload, String comment);
 }
diff --git a/src/main/java/connection/WebSocketConnection.java b/src/main/java/connection/WebSocketConnection.java
index b1f9ad9..c9bc0c9 100644
--- a/src/main/java/connection/WebSocketConnection.java
+++ b/src/main/java/connection/WebSocketConnection.java
@@ -53,6 +53,21 @@ public void queue(String payload)
         }
     }
 
+    @Override
+    public void queue(String payload, String comment)
+    {
+        if (isAttackRunning.get())
+        {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");
+            }
+        }
+    }
+
     public void sendMessage(String payload)
     {
         extensionWebSocket.sendTextMessage(payload);
diff --git a/src/main/resources/examples/BasicExampleWithComments.py b/src/main/resources/examples/BasicExampleWithComments.py
new file mode 100644
index 0000000..2646183
--- /dev/null
+++ b/src/main/resources/examples/BasicExampleWithComments.py
@@ -0,0 +1,12 @@
+def queue_websockets(base_websocket, payload):
+    connection1 = websocket_connection.create(base_websocket)
+
+    for i in range(10):
+        connection1.queue(payload, "foo")
+
+def handle_outgoing_message(websocket_message):
+    results_table.add(websocket_message)
+
+def handle_incoming_message(websocket_message):
+    websocket_message.setComment("bar")
+    results_table.add(websocket_message)
\ No newline at end of file