various prompt and UI tweaks

added support for new placeholders

README.md

@@ -65,6 +65,8 @@ After configuring the extension with the appropriate `API key`, `model`, and `pr
 
 `burpgpt` allows users to customise the `prompt` for traffic-based analysis by using a system of `placeholders`. We recommend including the maximum relevant information in the prompt. The following `placeholders` are directly handled by the extension and can be used to dynamically insert specific values into the prompt:
 
+- `{REQUEST}` - The scanned request.
+- `{RESPONSE}` - The scanned response.
 - `{IS_TRUNCATED_PROMPT}` - A `boolean` value that indicates whether the `prompt` has been truncated to fit within the `1024 character` limit imposed by most `GPT-3.5` models' `maxTokens` value. This value is programmatically set by the extenstion.
 - `{URL}` - The URL of the scanned request.
 - `{METHOD}` - The HTTP request method used in the scanned request.

lib/build.gradle

@@ -31,6 +31,7 @@ dependencies {
     implementation 'com.google.code.gson:gson:2.8.9' 
     implementation 'com.squareup.okhttp3:okhttp:4.9.3'
     implementation 'net.portswigger.burp.extensions:montoya-api:2023.3'
+    implementation 'org.apache.commons:commons-text:1.9'
 
     // Use JUnit Jupiter for testing.
     testImplementation 'net.portswigger.burp.extensions:montoya-api:2023.3'

lib/src/main/java/burp/MyBurpExtension.java

@@ -17,7 +17,7 @@
 public class MyBurpExtension implements BurpExtension, PropertyChangeListener {
 
     public static final String EXTENSION = "BurpGPT";
-    public static final Boolean DEBUG = false;
+    public static final Boolean DEBUG = true;
 
     private PropertyChangeSupport propertyChangeSupport;
     @Getter
@@ -30,23 +30,15 @@ public class MyBurpExtension implements BurpExtension, PropertyChangeListener {
     @Getter
     private String modelId = modelIds.get(0);
     @Getter
-    private String prompt = "Analyze the following HTTP request and response for potential vulnerabilities, "
-            + "including but not limited to OWASP top 10 vulnerabilities,such as SQL injection, XSS, CSRF, and more. "
-            + "Other common web application security threats should also be considered.\n"
-            + "Is Truncated Prompt: {IS_TRUNCATED_PROMPT}\n"
-            + "The request was made using the following information:\n"
-            + "URL: {URL}\n"
-            + "Request Method: {METHOD}\n"
-            + "Request Headers: {REQUEST_HEADERS}\n"
-            + "Request Body: {REQUEST_BODY}\n"
-            + "The response contains the following information:\n"
-            + "Response Headers: {RESPONSE_HEADERS}\n"
-            + "Response Body: {RESPONSE_BODY}\n"
-            + "Please describe any identified vulnerabilities and explain how they could be exploited.\n"
-            + "Each vulnerability should be formatted as follows:\n"
-            + "- Vulnerability Name: Brief description of vulnerability\n"
-            + "For example:\n"
-            + "- SQL Injection: The application is vulnerable to SQL injection attacks due to unsanitized user input in the search form.";;
+    String prompt = "Please analyze the following HTTP request and response for potential security vulnerabilities, "
+            + "specifically focusing on OWASP top 10 vulnerabilities such as SQL injection, XSS, CSRF, and other common web application security threats.\n\n"
+            + "Format your response as a bullet list with each point listing a vulnerability name and a brief description, in the format:\n"
+            + "- Vulnerability Name: Brief description of vulnerability\n\n"
+            + "Exclude irrelevant information.\n\n"
+            + "=== Request ===\n"
+            + "{REQUEST}\n\n"
+            + "=== Response ===\n"
+            + "{RESPONSE}\n";
 
     private GPTClient gptClient;
 

lib/src/main/java/burp/MyScanCheck.java

@@ -3,7 +3,9 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.stream.Collectors;
+
+import org.apache.commons.lang3.tuple.Pair;
+import org.apache.commons.text.StringEscapeUtils;
 
 import burp.api.montoya.http.message.HttpRequestResponse;
 import burp.api.montoya.logging.Logging;
@@ -14,6 +16,7 @@
 import burp.api.montoya.scanner.audit.issues.AuditIssue;
 import burp.api.montoya.scanner.audit.issues.AuditIssueConfidence;
 import burp.api.montoya.scanner.audit.issues.AuditIssueSeverity;
+import burpgpt.gpt.GPTRequest;
 import burpgpt.gpt.GPTResponse;
 import burpgpt.http.GPTClient;
 import lombok.Setter;
@@ -38,8 +41,8 @@ public AuditResult activeAudit(HttpRequestResponse httpRequestResponse, AuditIns
     @Override
     public AuditResult passiveAudit(HttpRequestResponse httpRequestResponse) {
         try {
-            GPTResponse gptResponse = gptClient.identifyVulnerabilities(httpRequestResponse);
-            List<AuditIssue> auditIssues = createAuditIssuesFromGPTResponse(gptResponse, httpRequestResponse);
+            Pair<GPTRequest, GPTResponse> gptResults = gptClient.identifyVulnerabilities(httpRequestResponse);
+            List<AuditIssue> auditIssues = createAuditIssuesFromGPTResponse(gptResults, httpRequestResponse);
             return AuditResult.auditResult(auditIssues);
         } catch (IOException e) {
             logging.raiseErrorEvent(e.getMessage());
@@ -49,23 +52,29 @@ public AuditResult passiveAudit(HttpRequestResponse httpRequestResponse) {
 
     @Override
     public ConsolidationAction consolidateIssues(AuditIssue newIssue, AuditIssue existingIssue) {
-        return existingIssue.name().equals(newIssue.name()) ? ConsolidationAction.KEEP_EXISTING
+        return newIssue.equals(existingIssue) ? ConsolidationAction.KEEP_EXISTING
                 : ConsolidationAction.KEEP_BOTH;
     }
 
-    private List<AuditIssue> createAuditIssuesFromGPTResponse(GPTResponse gptResponse,
+    private List<AuditIssue> createAuditIssuesFromGPTResponse(Pair<GPTRequest, GPTResponse> gptResults,
             HttpRequestResponse httpRequestResponse) {
         List<AuditIssue> auditIssues = new ArrayList<>();
-        String choicesText = "";
+        GPTRequest gptRequest = gptResults.getLeft();
+        GPTResponse gptResponse = gptResults.getRight();
 
         if (gptResponse.getChoices() != null) {
-            choicesText = gptResponse.getChoices().stream()
-                    .map(choice -> choice.getText())
-                    .collect(Collectors.joining("\n"));
+            String escapedPrompt = StringEscapeUtils.escapeHtml4(gptRequest.getPrompt().trim()).replace("\n", "<br />");
+            String issueBackground = String.format(
+                    "The following prompt was sent to the OpenAI %s GPT model and generate a response" +
+                            "based on the selected HTTP request and response:<br><br>%s",
+                    gptRequest.getModelId(), escapedPrompt);
+
+            String choiceText = gptResponse.getChoices().get(0).getText();
+            String issueDetail = StringEscapeUtils.escapeHtml4(choiceText.trim()).replace("\n", "<br />");
 
-            AuditIssue auditIssue = AuditIssue.auditIssue("GPT-generated vulnerability insights", choicesText,
+            AuditIssue auditIssue = AuditIssue.auditIssue("GPT-generated vulnerability insights", issueDetail,
                     null, httpRequestResponse.request().url(), AuditIssueSeverity.INFORMATION,
-                    AuditIssueConfidence.TENTATIVE, null, null,
+                    AuditIssueConfidence.TENTATIVE, issueBackground, null,
                     null, httpRequestResponse);
             auditIssues.add(auditIssue);
         }

lib/src/main/java/burpgpt/gpt/GPTRequest.java

@@ -8,32 +8,43 @@ public class GPTRequest {
     @Getter
     private String prompt;
     @Getter
-    private final int maxTokens;
+    private String modelId;
     @Getter
     private final int n;
+    @Getter
+    private final int maxTokens;
     private final String url;
     private final String method;
     private final String requestHeaders;
     private final String requestBody;
     private final String responseHeaders;
     private final String responseBody;
 
-    public GPTRequest(HttpRequest httpRequest, HttpResponse httpResponse, int n, int maxTokens) {
+    private final String request;
+    private final String response;
+
+    public GPTRequest(HttpRequest httpRequest, HttpResponse httpResponse, String modelId, int n, int maxTokens) {
         this.url = httpRequest.url();
         this.method = httpRequest.method();
         this.requestHeaders = httpRequest.headers().toString();
         this.requestBody = httpRequest.bodyToString();
         this.responseHeaders = httpResponse.headers().toString();
         this.responseBody = httpResponse.bodyToString();
 
+        this.request = httpRequest.toString();
+        this.response = httpResponse.toString();
+
+        this.modelId = modelId;
         this.n = n;
         this.maxTokens = maxTokens;
     }
 
     public void setPrompt(String prompt) {
-        String[] placeholders = { "{IS_TRUNCATED_PROMPT}", "{URL}", "{METHOD}", "{REQUEST_HEADERS}", "{REQUEST_BODY}",
+        String[] placeholders = { "{REQUEST}", "{RESPONSE}", "{IS_TRUNCATED_PROMPT}", "{URL}", "{METHOD}",
+                "{REQUEST_HEADERS}", "{REQUEST_BODY}",
                 "{RESPONSE_HEADERS}", "{RESPONSE_BODY}" };
-        String[] replacements = { Boolean.toString(prompt.length() > maxTokens), url, method, requestHeaders,
+        String[] replacements = { request, response, Boolean.toString(prompt.length() > maxTokens), url, method,
+                requestHeaders,
                 requestBody, responseHeaders, responseBody };
 
         for (int i = 0; i < placeholders.length; i++) {

lib/src/main/java/burpgpt/gpt/GPTResponse.java

@@ -2,54 +2,76 @@
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
+
+import com.google.gson.annotations.SerializedName;
 
 import lombok.Getter;
 
 @Getter
 public class GPTResponse {
     private List<Choice> choices;
-    private List<Vulnerability> vulnerabilities;
+    private String modelId;
+    private String id;
+    @SerializedName("created")
+    private long createdTimestamp;
+    @SerializedName("usage")
+    private Usage usage;
 
     public GPTResponse(List<Choice> choices) {
         this.choices = choices;
-        if (choices != null && !choices.isEmpty()) {
-            parseVulnerabilities();
-        }
     }
 
     @Getter
     public class Choice {
         private String text;
-    }
+        private int index;
+        private Object logprobs; // or use a specific class structure if needed
+        @SerializedName("finish_reason")
+        private String finishReason;
 
-    private void parseVulnerabilities() {
-        vulnerabilities = new ArrayList<>();
+        @Override
+        public String toString() {
+            return "Choice{" +
+                    "text='" + text + '\'' +
+                    ", index=" + index +
+                    ", logprobs=" + logprobs +
+                    ", finishReason='" + finishReason + '\'' +
+                    '}';
+        }
+    }
 
+    public List<String> getChoiceTexts() {
+        List<String> choiceTexts = new ArrayList<>();
         for (Choice choice : choices) {
-            String text = choice.getText();
-            if (text != null && !text.isBlank()) {
-                Pattern pattern = Pattern.compile("- ([^:]+): ([\\s\\S]+?)(?=- \\w|$)");
-                Matcher matcher = pattern.matcher(text);
-
-                while (matcher.find()) {
-                    String vulnerabilityName = matcher.group(1).trim();
-                    String vulnerabilityDescription = matcher.group(2).trim();
-                    vulnerabilities.add(new Vulnerability(vulnerabilityName, vulnerabilityDescription));
-                }
-            }
+            choiceTexts.add(choice.getText());
         }
+        return choiceTexts;
     }
 
     @Getter
-    public static class Vulnerability {
-        private String name;
-        private String description;
+    public static class Usage {
+        private long promptTokens;
+        private long completionTokens;
+        private long totalTokens;
 
-        public Vulnerability(String name, String description) {
-            this.name = name;
-            this.description = description;
+        @Override
+        public String toString() {
+            return "Usage{" +
+                    "promptTokens=" + promptTokens +
+                    ", completionTokens=" + completionTokens +
+                    ", totalTokens=" + totalTokens +
+                    '}';
         }
     }
-}
+
+    @Override
+    public String toString() {
+        return "GPTResponse{" +
+                "choices=" + choices +
+                ", modelId='" + modelId + '\'' +
+                ", id='" + id + '\'' +
+                ", createdTimestamp=" + createdTimestamp +
+                ", usage=" + usage +
+                '}';
+    }
+}

lib/src/main/java/burpgpt/http/GPTClient.java

@@ -3,6 +3,8 @@
 import java.io.IOException;
 import java.util.concurrent.TimeUnit;
 
+import org.apache.commons.lang3.tuple.Pair;
+
 import com.google.gson.Gson;
 import com.google.gson.JsonObject;
 
@@ -48,7 +50,7 @@ public void updateSettings(String newApiKey, String newModelId, String newPrompt
     this.prompt = newPrompt;
   }
 
-  public GPTResponse identifyVulnerabilities(HttpRequestResponse selectedMessage) throws IOException {
+  public Pair<GPTRequest, GPTResponse> identifyVulnerabilities(HttpRequestResponse selectedMessage) throws IOException {
     HttpRequest selectedRequest = selectedMessage.request();
     HttpResponse selectedResponse = selectedMessage.response();
 
@@ -74,8 +76,9 @@ public GPTResponse identifyVulnerabilities(HttpRequestResponse selectedMessage)
     // and receives a list of potential vulnerabilities in response.
     // TODO: Add a field to specify the maxTokens value
     try {
-      GPTRequest gptRequest = new GPTRequest(selectedRequest, selectedResponse, 1, 1024);
-      return getCompletions(gptRequest, apiKey, modelId, prompt);
+      GPTRequest gptRequest = new GPTRequest(selectedRequest, selectedResponse, modelId, 1, 1024);
+      GPTResponse gptResponse = getCompletions(gptRequest, apiKey, modelId, prompt);
+      return Pair.of(gptRequest, gptResponse);
     } catch (IOException e) {
       throw e;
     }