diff --git a/plugins/burp-suite/cyberstrikeai-burp-extension/dist/cyberstrikeai-burp-extension.jar b/plugins/burp-suite/cyberstrikeai-burp-extension/dist/cyberstrikeai-burp-extension.jar index 5f64ac3e..789f49da 100644 Binary files a/plugins/burp-suite/cyberstrikeai-burp-extension/dist/cyberstrikeai-burp-extension.jar and b/plugins/burp-suite/cyberstrikeai-burp-extension/dist/cyberstrikeai-burp-extension.jar differ diff --git a/plugins/burp-suite/cyberstrikeai-burp-extension/src/main/java/burp/BurpExtender.java b/plugins/burp-suite/cyberstrikeai-burp-extension/src/main/java/burp/BurpExtender.java index 40c011ca..1f40da51 100644 --- a/plugins/burp-suite/cyberstrikeai-burp-extension/src/main/java/burp/BurpExtender.java +++ b/plugins/burp-suite/cyberstrikeai-burp-extension/src/main/java/burp/BurpExtender.java @@ -1,6 +1,7 @@ package burp; import javax.swing.*; +import java.awt.*; import java.util.ArrayList; import java.util.List; @@ -10,6 +11,7 @@ public class BurpExtender implements IBurpExtender, IContextMenuFactory { private CyberStrikeAITab tab; private final CyberStrikeAIClient client = new CyberStrikeAIClient(); + private String lastInstruction = HttpMessageFormatter.defaultInstruction(); @Override public void registerExtenderCallbacks(IBurpExtenderCallbacks callbacks) { @@ -36,111 +38,149 @@ public class BurpExtender implements IBurpExtender, IContextMenuFactory { if (selected == null || selected.length == 0) { return; } - - CyberStrikeAIClient.Config cfg = tab.currentConfig(); - String token = tab.getToken(); - if (token == null || token.trim().isEmpty()) { - JOptionPane.showMessageDialog(tab.getUiComponent(), - "Please click Validate first to obtain a token.", - "CyberStrikeAI", JOptionPane.WARNING_MESSAGE); - return; - } - - String prompt = HttpMessageFormatter.toPrompt(helpers, selected[0]); - String title = HttpMessageFormatter.getRequestTitle(helpers, selected[0]); - String agentModeStr = (cfg.agentMode == CyberStrikeAIClient.AgentMode.MULTI) ? "Multi Agent" : "Single Agent"; - String runId = tab.startNewRun(title, agentModeStr, selected[0]); - tab.appendProgressToRun(runId, "\n[server] " + cfg.baseUrl + "\n\n"); - - client.streamTest(cfg, token, prompt, new CyberStrikeAIClient.StreamListener() { - @Override - public void onEvent(String type, String message, String rawJson) { - if (type == null) type = ""; - switch (type) { - case "response_delta": - case "eino_agent_reply_stream_delta": - // delta chunk (content only) - tab.appendFinalToRun(runId, message); - break; - case "response": - // final response (full) - tab.appendFinalToRun(runId, "\n\n--- Final Response ---\n"); - tab.appendFinalToRun(runId, message); - tab.setFinalResponse(runId, message); - break; - case "progress": - tab.appendProgressToRun(runId, "\n[progress] " + message + "\n"); - tab.setRunStatus(runId, "running"); - break; - case "cancelled": - tab.appendProgressToRun(runId, "\n[cancelled] " + message + "\n"); - tab.setRunStatus(runId, "cancelled"); - break; - case "error": - tab.appendProgressToRun(runId, "\n[error] " + message + "\n"); - tab.setRunStatus(runId, "error"); - break; - case "thinking_stream_start": - if (tab.isShowDebugEvents()) { - tab.resetThinkingStream(runId); - } - break; - case "thinking_stream_delta": - case "tool_call": - case "tool_result": - case "tool_result_delta": - // debug; hide by default - if (tab.isShowDebugEvents() && message != null && !message.isEmpty()) { - if ("thinking_stream_delta".equals(type)) { - tab.appendThinkingDelta(runId, message); - } else { - tab.appendProgressToRun(runId, "\n[" + type + "] " + message + "\n"); - } - } - break; - case "conversation": - // Capture conversationId for stop/cancel. - if (rawJson != null) { - String convId = SimpleJson.extractStringField(rawJson, "conversationId"); - if (convId != null && !convId.trim().isEmpty()) { - tab.setRunConversationId(runId, convId); - } - } - if (tab.isShowDebugEvents() && message != null && !message.isEmpty()) { - tab.appendProgressToRun(runId, "\n[" + type + "] " + message + "\n"); - } - break; - case "done": - // handled in onDone too - break; - default: - if (tab.isShowDebugEvents() && message != null && !message.isEmpty()) { - tab.appendProgressToRun(runId, "\n[" + type + "] " + message + "\n"); - } - break; - } - } - - @Override - public void onError(String message, Exception e) { - tab.appendProgressToRun(runId, "\n[error] " + message + "\n"); - tab.setRunStatus(runId, "error"); - callbacks.printError("CyberStrikeAI stream error: " + message); - if (e != null) { - callbacks.printError(e.toString()); - } - } - - @Override - public void onDone() { - tab.appendProgressToRun(runId, "\n\n[done]\n"); - tab.setRunStatus(runId, "done"); - } - }); + sendMessage(selected[0]); }); items.add(sendItem); return items; } + + private void sendMessage(IHttpRequestResponse msg) { + if (msg == null) return; + CyberStrikeAIClient.Config cfg = tab.currentConfig(); + String token = tab.getToken(); + if (token == null || token.trim().isEmpty()) { + JOptionPane.showMessageDialog(tab.getUiComponent(), + "Please click Validate first to obtain a token.", + "CyberStrikeAI", JOptionPane.WARNING_MESSAGE); + return; + } + + String instruction = showInstructionEditor(tab.getUiComponent(), lastInstruction); + if (instruction == null) { + return; + } + lastInstruction = instruction; + + String prompt = HttpMessageFormatter.toPrompt(helpers, msg, instruction); + String title = HttpMessageFormatter.getRequestTitle(helpers, msg); + String agentModeStr = (cfg.agentMode == CyberStrikeAIClient.AgentMode.MULTI) ? "Multi Agent" : "Single Agent"; + String runId = tab.startNewRun(title, agentModeStr, msg); + tab.appendProgressToRun(runId, "\n[server] " + cfg.baseUrl + "\n\n"); + + client.streamTest(cfg, token, prompt, new CyberStrikeAIClient.StreamListener() { + @Override + public void onEvent(String type, String message, String rawJson) { + if (type == null) type = ""; + switch (type) { + case "response_delta": + case "eino_agent_reply_stream_delta": + tab.appendFinalToRun(runId, message); + break; + case "response": + tab.appendFinalToRun(runId, "\n\n--- Final Response ---\n"); + tab.appendFinalToRun(runId, message); + tab.setFinalResponse(runId, message); + break; + case "progress": + tab.appendProgressToRun(runId, "\n[progress] " + message + "\n"); + tab.setRunStatus(runId, "running"); + break; + case "cancelled": + tab.appendProgressToRun(runId, "\n[cancelled] " + message + "\n"); + tab.setRunStatus(runId, "cancelled"); + break; + case "error": + tab.appendProgressToRun(runId, "\n[error] " + message + "\n"); + tab.setRunStatus(runId, "error"); + break; + case "thinking_stream_start": + if (tab.isShowDebugEvents()) { + tab.resetThinkingStream(runId); + } + break; + case "thinking_stream_delta": + case "tool_call": + case "tool_result": + case "tool_result_delta": + if (tab.isShowDebugEvents() && message != null && !message.isEmpty()) { + if ("thinking_stream_delta".equals(type)) { + tab.appendThinkingDelta(runId, message); + } else { + tab.appendProgressToRun(runId, "\n[" + type + "] " + message + "\n"); + } + } + break; + case "conversation": + if (rawJson != null) { + String convId = SimpleJson.extractStringField(rawJson, "conversationId"); + if (convId != null && !convId.trim().isEmpty()) { + tab.setRunConversationId(runId, convId); + } + } + if (tab.isShowDebugEvents() && message != null && !message.isEmpty()) { + tab.appendProgressToRun(runId, "\n[" + type + "] " + message + "\n"); + } + break; + case "done": + break; + default: + if (tab.isShowDebugEvents() && message != null && !message.isEmpty()) { + tab.appendProgressToRun(runId, "\n[" + type + "] " + message + "\n"); + } + break; + } + } + + @Override + public void onError(String message, Exception e) { + tab.appendProgressToRun(runId, "\n[error] " + message + "\n"); + tab.setRunStatus(runId, "error"); + callbacks.printError("CyberStrikeAI stream error: " + message); + if (e != null) { + callbacks.printError(e.toString()); + } + } + + @Override + public void onDone() { + tab.appendProgressToRun(runId, "\n\n[done]\n"); + tab.setRunStatus(runId, "done"); + } + }); + } + + private static String showInstructionEditor(Component parent, String initialValue) { + JTextArea editor = new JTextArea( + initialValue == null || initialValue.trim().isEmpty() + ? HttpMessageFormatter.defaultInstruction() + : initialValue, + 6, + 70 + ); + editor.setLineWrap(true); + editor.setWrapStyleWord(true); + editor.setFont(new Font(Font.SANS_SERIF, Font.PLAIN, 13)); + + JPanel panel = new JPanel(new BorderLayout(0, 8)); + panel.add(new JLabel("Edit instruction before sending:"), BorderLayout.NORTH); + panel.add(new JScrollPane(editor), BorderLayout.CENTER); + + int result = JOptionPane.showConfirmDialog( + parent, + panel, + "Customize Prompt Instruction", + JOptionPane.OK_CANCEL_OPTION, + JOptionPane.PLAIN_MESSAGE + ); + if (result != JOptionPane.OK_OPTION) { + return null; + } + String value = editor.getText(); + if (value == null || value.trim().isEmpty()) { + return HttpMessageFormatter.defaultInstruction(); + } + return value.trim(); + } } diff --git a/plugins/burp-suite/cyberstrikeai-burp-extension/src/main/java/burp/HttpMessageFormatter.java b/plugins/burp-suite/cyberstrikeai-burp-extension/src/main/java/burp/HttpMessageFormatter.java index 792465b0..41bc69c2 100644 --- a/plugins/burp-suite/cyberstrikeai-burp-extension/src/main/java/burp/HttpMessageFormatter.java +++ b/plugins/burp-suite/cyberstrikeai-burp-extension/src/main/java/burp/HttpMessageFormatter.java @@ -5,6 +5,8 @@ import java.util.List; final class HttpMessageFormatter { private HttpMessageFormatter() {} + private static final String DEFAULT_INSTRUCTION = + "针对该流量做web渗透测试，并输出测试结果，要求：只针对该接口流量做测试，切勿拓展其他接口"; static String getRequestTitle(IExtensionHelpers helpers, IHttpRequestResponse msg) { IRequestInfo reqInfo = helpers.analyzeRequest(msg); @@ -22,7 +24,15 @@ final class HttpMessageFormatter { return method + " " + host + shortPath + q; } + static String defaultInstruction() { + return DEFAULT_INSTRUCTION; + } + static String toPrompt(IExtensionHelpers helpers, IHttpRequestResponse msg) { + return toPrompt(helpers, msg, DEFAULT_INSTRUCTION); + } + + static String toPrompt(IExtensionHelpers helpers, IHttpRequestResponse msg, String instruction) { IRequestInfo reqInfo = helpers.analyzeRequest(msg); String method = reqInfo.getMethod(); String url = reqInfo.getUrl() != null ? reqInfo.getUrl().toString() : "(unknown)"; @@ -53,8 +63,12 @@ final class HttpMessageFormatter { + respBody; } + String prefix = (instruction == null || instruction.trim().isEmpty()) + ? DEFAULT_INSTRUCTION + : instruction.trim(); + return "" - + "针对该流量做web渗透测试，并输出测试结果，要求：只针对该接口流量做测试，切勿拓展其他接口\n\n" + + prefix + "\n\n" + "[Target]\n" + method + " " + url + "\n\n" + "[Request]\n" diff --git a/plugins/burp-suite/cyberstrikeai-burp-extension/target/classes/burp/BurpExtender$1.class b/plugins/burp-suite/cyberstrikeai-burp-extension/target/classes/burp/BurpExtender$1.class index d3f3b4a2..4921cbb9 100644 Binary files a/plugins/burp-suite/cyberstrikeai-burp-extension/target/classes/burp/BurpExtender$1.class and b/plugins/burp-suite/cyberstrikeai-burp-extension/target/classes/burp/BurpExtender$1.class differ diff --git a/plugins/burp-suite/cyberstrikeai-burp-extension/target/classes/burp/BurpExtender.class b/plugins/burp-suite/cyberstrikeai-burp-extension/target/classes/burp/BurpExtender.class index 396d2c69..a271921f 100644 Binary files a/plugins/burp-suite/cyberstrikeai-burp-extension/target/classes/burp/BurpExtender.class and b/plugins/burp-suite/cyberstrikeai-burp-extension/target/classes/burp/BurpExtender.class differ diff --git a/plugins/burp-suite/cyberstrikeai-burp-extension/target/classes/burp/HttpMessageFormatter.class b/plugins/burp-suite/cyberstrikeai-burp-extension/target/classes/burp/HttpMessageFormatter.class index e4a781e5..0ace303f 100644 Binary files a/plugins/burp-suite/cyberstrikeai-burp-extension/target/classes/burp/HttpMessageFormatter.class and b/plugins/burp-suite/cyberstrikeai-burp-extension/target/classes/burp/HttpMessageFormatter.class differ diff --git a/plugins/burp-suite/cyberstrikeai-burp-extension/target/cyberstrikeai-burp-extension-1.0.0.jar b/plugins/burp-suite/cyberstrikeai-burp-extension/target/cyberstrikeai-burp-extension-1.0.0.jar index 5f64ac3e..789f49da 100644 Binary files a/plugins/burp-suite/cyberstrikeai-burp-extension/target/cyberstrikeai-burp-extension-1.0.0.jar and b/plugins/burp-suite/cyberstrikeai-burp-extension/target/cyberstrikeai-burp-extension-1.0.0.jar differ