diff --git a/README.md b/README.md index 2bcadbc..9ecc777 100644 --- a/README.md +++ b/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. diff --git a/lib/build.gradle b/lib/build.gradle index 4c0aa82..fa62b7d 100644 --- a/lib/build.gradle +++ b/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' diff --git a/lib/src/main/java/burp/MyBurpExtension.java b/lib/src/main/java/burp/MyBurpExtension.java index a698006..4b6e3d6 100644 --- a/lib/src/main/java/burp/MyBurpExtension.java +++ b/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; diff --git a/lib/src/main/java/burp/MyScanCheck.java b/lib/src/main/java/burp/MyScanCheck.java index 203e044..9703c73 100644 --- a/lib/src/main/java/burp/MyScanCheck.java +++ b/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 auditIssues = createAuditIssuesFromGPTResponse(gptResponse, httpRequestResponse); + Pair gptResults = gptClient.identifyVulnerabilities(httpRequestResponse); + List 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 createAuditIssuesFromGPTResponse(GPTResponse gptResponse, + private List createAuditIssuesFromGPTResponse(Pair gptResults, HttpRequestResponse httpRequestResponse) { List 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", "
"); + 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:

%s", + gptRequest.getModelId(), escapedPrompt); + + String choiceText = gptResponse.getChoices().get(0).getText(); + String issueDetail = StringEscapeUtils.escapeHtml4(choiceText.trim()).replace("\n", "
"); - 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); } diff --git a/lib/src/main/java/burpgpt/gpt/GPTRequest.java b/lib/src/main/java/burpgpt/gpt/GPTRequest.java index 964174e..a609bc1 100644 --- a/lib/src/main/java/burpgpt/gpt/GPTRequest.java +++ b/lib/src/main/java/burpgpt/gpt/GPTRequest.java @@ -8,9 +8,11 @@ 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; @@ -18,7 +20,10 @@ public class GPTRequest { 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(); @@ -26,14 +31,20 @@ public GPTRequest(HttpRequest httpRequest, HttpResponse httpResponse, int n, int 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++) { diff --git a/lib/src/main/java/burpgpt/gpt/GPTResponse.java b/lib/src/main/java/burpgpt/gpt/GPTResponse.java index 4269aae..c58c6b8 100644 --- a/lib/src/main/java/burpgpt/gpt/GPTResponse.java +++ b/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 choices; - private List vulnerabilities; + private String modelId; + private String id; + @SerializedName("created") + private long createdTimestamp; + @SerializedName("usage") + private Usage usage; public GPTResponse(List 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 getChoiceTexts() { + List 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 + + '}'; } } -} \ No newline at end of file + + @Override + public String toString() { + return "GPTResponse{" + + "choices=" + choices + + ", modelId='" + modelId + '\'' + + ", id='" + id + '\'' + + ", createdTimestamp=" + createdTimestamp + + ", usage=" + usage + + '}'; + } +} diff --git a/lib/src/main/java/burpgpt/http/GPTClient.java b/lib/src/main/java/burpgpt/http/GPTClient.java index da67cc8..e759d78 100644 --- a/lib/src/main/java/burpgpt/http/GPTClient.java +++ b/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 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; }