Changeable AI provider (with OpenAI-compatible alternative)

.gitignore

@@ -9,4 +9,6 @@ out/
 .project/
 /build/
 /.gradle/
-libs/montoya-api-2025.3.6.jar
+libs/montoya-api-2025.3.6.jar
+#other
+.DS_Store

src/main/java/burp/shadow/repeater/ai/AI.java

@@ -3,8 +3,6 @@
 import burp.shadow.repeater.ShadowRepeaterExtension;
 import burp.shadow.repeater.settings.InvalidTypeSettingException;
 import burp.shadow.repeater.settings.UnregisteredSettingException;
-import burp.api.montoya.ai.chat.Message;
-import burp.api.montoya.ai.chat.PromptOptions;
 import burp.api.montoya.ai.chat.PromptResponse;
 
 import java.security.MessageDigest;
@@ -71,7 +69,7 @@ public String execute() {
             api.logging().logToOutput("System Prompt:" + this.systemMessage + "\n\n");
             api.logging().logToOutput("Prompt:" + this.prompt + "\n\n");
         }
-        PromptResponse response = api.ai().prompt().execute(PromptOptions.promptOptions().withTemperature(this.temperature), Message.systemMessage(this.systemMessage), Message.userMessage(this.prompt));
+        PromptResponse response = AIProvider.acquire().execute(this.temperature, this.systemMessage, this.prompt);
         if(debugAi) {
             api.logging().logToOutput("AI Response:" + response.content() + "\n\n");
         }

src/main/java/burp/shadow/repeater/ai/AIExecutor.java

@@ -0,0 +1,7 @@
+package burp.shadow.repeater.ai;
+
+import burp.api.montoya.ai.chat.PromptResponse;
+
+public interface AIExecutor {
+    public PromptResponse execute(double temperature, String systemMessage, String userMessage);
+}

src/main/java/burp/shadow/repeater/ai/AIProvider.java

@@ -0,0 +1,23 @@
+package burp.shadow.repeater.ai;
+
+import burp.shadow.repeater.ShadowRepeaterExtension;
+import burp.shadow.repeater.ai.executor.BurpAIExecutor;
+import burp.shadow.repeater.ai.executor.OpenAIExecutor;
+import burp.shadow.repeater.settings.InvalidTypeSettingException;
+import burp.shadow.repeater.settings.UnregisteredSettingException;
+
+public class AIProvider {
+    public static AIExecutor acquire() {
+        AIProviderType aiProvider;
+        try {
+            aiProvider = AIProviderType.valueOf(ShadowRepeaterExtension.generalSettings.getStringEnum("aiProvider"));
+        } catch (UnregisteredSettingException | InvalidTypeSettingException e) {
+            aiProvider = AIProviderType.BurpAI;
+        }
+
+        return switch (aiProvider) {
+            case AIProviderType.OpenAI -> new OpenAIExecutor();
+            default -> new BurpAIExecutor(); // Return Burp AI by default
+        };
+    };
+}

src/main/java/burp/shadow/repeater/ai/AIProviderType.java

@@ -0,0 +1,15 @@
+package burp.shadow.repeater.ai;
+
+public enum AIProviderType {
+    BurpAI("BurpAI"), OpenAI("OpenAI");
+
+    private final String value;
+
+    AIProviderType(String value) {
+        this.value = value;
+    }
+
+    public String value() {
+        return value;
+    }
+}

src/main/java/burp/shadow/repeater/ai/executor/BurpAIExecutor.java

@@ -0,0 +1,15 @@
+package burp.shadow.repeater.ai.executor;
+
+import burp.api.montoya.ai.chat.Message;
+import burp.api.montoya.ai.chat.PromptOptions;
+import burp.api.montoya.ai.chat.PromptResponse;
+import burp.shadow.repeater.ai.AIExecutor;
+
+import static burp.shadow.repeater.ShadowRepeaterExtension.api;
+
+public class BurpAIExecutor implements AIExecutor {
+    @Override
+    public PromptResponse execute(double temperature, String systemMessage, String userMessage) {
+        return api.ai().prompt().execute(PromptOptions.promptOptions().withTemperature(temperature), Message.systemMessage(systemMessage), Message.userMessage(userMessage));
+    }
+}

src/main/java/burp/shadow/repeater/ai/executor/OpenAIExecutor.java

@@ -0,0 +1,97 @@
+package burp.shadow.repeater.ai.executor;
+
+import burp.api.montoya.ai.chat.PromptResponse;
+import burp.shadow.repeater.ai.AIExecutor;
+
+import org.json.JSONArray;
+import org.json.JSONObject;
+
+import static burp.shadow.repeater.ShadowRepeaterExtension.api;
+import static burp.shadow.repeater.ShadowRepeaterExtension.generalSettings;
+
+class OpenAIResponse implements PromptResponse {
+    private String content;
+
+    OpenAIResponse(String content) {
+        this.content = content;
+    }
+
+    public String content() {
+        return content;
+    }
+}
+
+class Payload extends JSONObject {
+    Payload(double temperature, String systemMessage, String userMessage, String model) {
+        put("model", (model != null && model.length() > 0) ? model : OpenAIExecutor.DEFAULT_MODEL);
+        put("temperature", temperature);
+
+        JSONArray messages = new JSONArray();
+
+        JSONObject message = new JSONObject();
+        message.put("role", "user");
+        message.put("content", userMessage);
+        messages.put(message);
+
+        message = new JSONObject();
+        message.put("role", "system");
+        message.put("content", systemMessage);
+        messages.put(message);
+
+        put("messages", messages);
+    }
+}
+
+public class OpenAIExecutor implements AIExecutor {
+    public static String DEFAULT_ENDPOINT = "https://api.openai.com/v1/responses";
+    public static String DEFAULT_MODEL = "o4-mini";
+
+    @Override
+    public PromptResponse execute(double temperature, String systemMessage, String userMessage) {
+        try {
+            String apiKey = generalSettings.getString("openAiApiKey");
+            String endpoint = generalSettings.getString("openAiEndpoint");
+            String model = generalSettings.getString("openAiModelName");
+
+            Payload payload = new Payload(temperature, systemMessage, userMessage, model);
+
+            java.net.http.HttpClient client = java.net.http.HttpClient.newHttpClient();
+            java.net.http.HttpRequest request;
+            java.net.http.HttpResponse<String> response;
+
+            request = java.net.http.HttpRequest.newBuilder()
+                .uri(java.net.URI.create(endpoint != null && endpoint.length() > 0 ? endpoint : DEFAULT_ENDPOINT))
+                .header("Content-Type", "application/json")
+                .header("Authorization", "Bearer " + apiKey)
+                .POST(java.net.http.HttpRequest.BodyPublishers.ofString(payload.toString()))
+                .build();
+
+            response = client.send(request, java.net.http.HttpResponse.BodyHandlers.ofString());
+
+            int status = response.statusCode();
+            String body = response.body();
+
+            if (status != 200) {
+                throw new Exception("Received response: " + status + " body: " + body);
+            }
+
+            // Parse the assistant content
+            JSONObject jsonResponse = new JSONObject(body);
+            JSONArray choices = jsonResponse.getJSONArray("choices");
+            String result = null;
+            for (int i = 0; i < choices.length(); i++) {
+                JSONObject choice = choices.getJSONObject(i);
+                JSONObject messageObj = choice.getJSONObject("message");
+                if ("assistant".equals(messageObj.optString("role"))) {
+                    result = messageObj.optString("content");
+                    break;
+                }
+            }
+
+            return new OpenAIResponse(result);
+        } catch (Exception e) {
+            api.logging().logToError(e);
+            return new OpenAIResponse("");
+        }
+    }
+}

src/main/java/burp/shadow/repeater/settings/Settings.java

@@ -4,15 +4,20 @@
 import burp.shadow.repeater.ShadowRepeaterExtension;
 import burp.shadow.repeater.utils.GridbagUtils;
 import burp.shadow.repeater.utils.Utils;
+
+import org.json.JSONArray;
 import org.json.JSONObject;
 
 import javax.swing.*;
 import javax.swing.border.EmptyBorder;
 import javax.swing.event.DocumentEvent;
 import javax.swing.event.DocumentListener;
 import java.awt.*;
+import java.awt.event.ActionEvent;
+import java.awt.event.ActionListener;
 import java.awt.event.ItemEvent;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
@@ -38,7 +43,7 @@ public static void showSettingsWindow() {
         settingsWindow.getContentPane().removeAll();
         settingsWindow.setTitle(extensionName + " settings");
         settingsWindow.setResizable(false);
-        settingsWindow.setPreferredSize(new Dimension(800, 600));
+        settingsWindow.setPreferredSize(new Dimension(800, 700));
         Container pane = settingsWindow.getContentPane();
         try {
             Map<String, Integer> columns = new HashMap<>();
@@ -70,7 +75,7 @@ public static void showSettingsWindow() {
     }
 
     public enum SettingType  {
-            Boolean, String, Password, Integer
+        Boolean, String, Password, Integer, StringEnum
     }
 
     public Settings(String settingsName, IBurpExtenderCallbacks callbacks) {
@@ -175,6 +180,23 @@ public void registerIntegerSetting(String name, int defaultValue, String descrip
         this.settings.put(name, setting);
         this.defaults.put(name, setting);
     }
+
+    public void registerStringEnumSetting(String name, String defaultValue, String description, String category, String[] values) {
+        if(!Arrays.asList(values).contains(defaultValue)){
+            throw new RuntimeException("Default value " + defaultValue + " is not present in possible values " + String.join(",", values));
+        }
+        addCategory(category, name);
+        JSONObject setting = this.settings.has(name) ? (JSONObject) this.settings.get(name) : new JSONObject();
+        setting.put("description", description);
+        setting.put("value", defaultValue);
+        setting.put("type", "StringEnum");
+        setting.put("category", category);
+        setting.put("values", new JSONArray(values));
+
+        this.settings.put(name, setting);
+        this.defaults.put(name, setting);
+    }
+
     public void load() {
         String json = callbacks.loadExtensionSetting(this.settingsName);
         if(json == null) {
@@ -228,6 +250,19 @@ public int getInteger(String name) throws UnregisteredSettingException, InvalidT
         throw new InvalidTypeSettingException("The setting " + name + " expects a int");
     }
 
+    public String getStringEnum(String name) throws UnregisteredSettingException, InvalidTypeSettingException {
+        JSONObject setting = this.getSetting(name);
+        String type = setting.getString("type");
+        if(SettingType.StringEnum.name().equals(type)) {
+            if(setting.has("value")) {
+                return setting.getString("value");
+            } else {
+                return this.defaults.getJSONObject(name).getString("value");
+            }
+        }
+        throw new InvalidTypeSettingException("The setting " + name + " expects a string enum");
+    }
+
     public void setString(String name, String value) throws UnregisteredSettingException, InvalidTypeSettingException {
         JSONObject setting = this.getSetting(name);
         String type = setting.getString("type");
@@ -260,6 +295,24 @@ public void setBoolean(String name, boolean value) throws UnregisteredSettingExc
         throw new InvalidTypeSettingException("Error setting " + name + " expects an boolean");
     }
 
+    public void setStringEnum(String name, String value) throws UnregisteredSettingException, InvalidTypeSettingException {
+        JSONObject setting = this.getSetting(name);
+        String type = setting.getString("type");
+
+        java.util.List<String> values = new ArrayList<>();
+        for(Object o : setting.getJSONArray("values")) {
+            values.add((String) o);
+        }
+
+        if(SettingType.StringEnum.name().equals(type) && values.contains(value)) {
+            setting.put("value", value);
+            isModified = true;
+            return;
+        }
+        throw new InvalidTypeSettingException("Error setting " + name
+                + " expects a string enum and the value must be one of " + String.join(",", values));
+    }
+
     private JSONObject getSetting(String name) throws UnregisteredSettingException {
         if(!this.settings.has(name) && !this.defaults.has(name)) {
             throw new UnregisteredSettingException(name +" has not been registered.");
@@ -293,6 +346,16 @@ private void updateBoolean(String name, boolean checked) {
         }
     }
 
+    private void updateStringEnum(String name, JComboBox<String> comboBox, JSONObject currentSetting) {
+        try {
+            String value = (String) comboBox.getSelectedItem();
+            this.setStringEnum(name, value);
+        } catch (UnregisteredSettingException | InvalidTypeSettingException ex) {
+            callbacks.printError(ex.toString());
+            throw new RuntimeException(ex);
+        }
+    }
+
     private void resetSettings() {
         Iterator<String> keys = this.settings.keys();
         while(keys.hasNext()) {
@@ -425,6 +488,32 @@ public void changedUpdate(DocumentEvent e) {
                         categoryContainer.add(label, addMarginToGbc(createConstraints(0, componentRow, 1, GridBagConstraints.BOTH, 1, 0, spacing, spacing, GridBagConstraints.WEST), 0, 5, 0,0));
                         categoryContainer.add(checkBox, createConstraints(1, componentRow, 1, GridBagConstraints.EAST, 0, 0, spacing, spacing, GridBagConstraints.EAST));
                     }
+                    case "StringEnum" -> {
+                        JLabel label = new JLabel(currentSetting.getString("description"));
+                        label.setPreferredSize(new Dimension(componentWidth, componentHeight));
+                        JComboBox<String> comboBox = new JComboBox<>();
+                        for (Object o : currentSetting.getJSONArray("values")) {
+                            comboBox.addItem((String) o);
+                        }
+                        comboBox.setEditable(false);
+                        comboBox.setSelectedItem(this.getStringEnum(name));
+                        comboBox.setPreferredSize(new Dimension(componentWidth, componentHeight));
+                        categoryContainer.add(label, addMarginToGbc(createConstraints(0, componentRow, 1,
+                                GridBagConstraints.BOTH, 1, 0, spacing, spacing, GridBagConstraints.WEST), 0, 5, 0, 0));
+                        categoryContainer.add(new JLabel(), createConstraints(1, componentRow, 1,
+                                GridBagConstraints.BOTH, 1, 0, spacing, spacing, GridBagConstraints.WEST));
+                        componentRow++;
+                        categoryContainer.add(comboBox,
+                                createConstraints(0, componentRow, 2, GridBagConstraints.BOTH, 1,
+                                        0, spacing, spacing, GridBagConstraints.WEST));
+
+                        comboBox.addActionListener(new ActionListener() {
+                            @Override
+                            public void actionPerformed(ActionEvent e) {
+                                updateStringEnum(name, comboBox, currentSetting);
+                            }
+                        });
+                    }
                     default -> {
                         throw new InvalidTypeSettingException("Unexpected type");
                     }

src/main/java/burp/shadow/repeater/utils/Utils.java

@@ -2,6 +2,8 @@
 
 import burp.api.montoya.http.message.params.ParsedHttpParameter;
 import burp.shadow.repeater.ShadowRepeaterExtension;
+import burp.shadow.repeater.ai.AIProviderType;
+import burp.shadow.repeater.ai.executor.OpenAIExecutor;
 import burp.shadow.repeater.settings.Settings;
 import burp.api.montoya.http.message.HttpRequestResponse;
 import burp.api.montoya.http.message.params.HttpParameter;
@@ -22,7 +24,6 @@
 import java.util.stream.Collectors;
 
 import static burp.shadow.repeater.ShadowRepeaterExtension.*;
-import static burp.shadow.repeater.ShadowRepeaterExtension.responseHistory;
 
 public class Utils {
     public static HttpRequest modifyRequest(HttpRequest req, String type, String name, String value) {
@@ -111,6 +112,10 @@ public static void registerGeneralSettings(Settings settings) {
         settings.registerBooleanSetting("debugOutput", false, "Print debug output", "General", null);
         settings.registerBooleanSetting("debugAi", false, "Debug AI requests/responses", "AI", null);
 
+        settings.registerStringEnumSetting("aiProvider", AIProviderType.BurpAI.value(), "AI provider", "AI", java.util.Arrays.stream(AIProviderType.values()).map(AIProviderType::value).toArray(String[]::new));
+        settings.registerPasswordSetting("openAiApiKey", "", "API key for OpenAI compatible APIs", "AI");
+        settings.registerStringSetting("openAiEndpoint", OpenAIExecutor.DEFAULT_ENDPOINT, "API endpoint URL for OpenAI compatible AI", "AI");
+        settings.registerStringSetting("openAiModelName", OpenAIExecutor.DEFAULT_MODEL, "The model used to generate the response", "AI");
     }
 
     public static void openUrl(String url) {