Add files via upload

This commit is contained in:
公明
2026-04-09 21:34:26 +08:00
committed by GitHub
parent cd30953a84
commit 9e34b480e7
7 changed files with 156 additions and 102 deletions
@@ -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();
}
}
@@ -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"