Skip to content

Commit 306f780

Browse files
przemyslvwprzemyslvw
committed
feat(ai): unify AiClient threading and timeouts
1 parent a763c96 commit 306f780

6 files changed

Lines changed: 171 additions & 145 deletions

File tree

‎src/main/java/burp/ReconMasterPro.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ public class ReconMasterPro implements BurpExtension {
2020
private CorsHunter corsHunter;
2121
private GraphQLExtractor graphqlExtractor;
2222
private CloudAssetsAggregator cloudAggregator;
23+
private AiAssistantPanel aiAssistantPanel;
2324

2425
@Override
2526
public void initialize(MontoyaApi api) {
@@ -98,7 +99,7 @@ public void initialize(MontoyaApi api) {
9899
);
99100
ReportPanel reportPanel = new ReportPanel(reportGenerator);
100101

101-
AiAssistantPanel aiAssistantPanel = new AiAssistantPanel(
102+
aiAssistantPanel = new AiAssistantPanel(
102103
settings,
103104
endpointsPanel::getEndpoints,
104105
techPanel::getTechnologies,
@@ -133,6 +134,7 @@ public void initialize(MontoyaApi api) {
133134
if (corsHunter != null) corsHunter.shutdown();
134135
if (graphqlExtractor != null) graphqlExtractor.shutdown();
135136
if (cloudAggregator != null) cloudAggregator.shutdown();
137+
if (aiAssistantPanel != null) aiAssistantPanel.shutdown();
136138
if (db != null) db.close();
137139
} catch (Exception ignored) {}
138140
});

‎src/main/java/burp/modules/AiClient.java‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,17 @@
1515
import java.time.Duration;
1616
import java.util.HashMap;
1717
import java.util.Map;
18+
import java.util.concurrent.ExecutorService;
19+
import java.util.concurrent.Executors;
1820

1921
public class AiClient {
2022

23+
private static final int CONNECT_TIMEOUT_SECONDS = 30;
24+
private static final int REQUEST_TIMEOUT_SECONDS = 90;
25+
2126
private final SettingsManager settings;
2227
private final HttpClient httpClient;
28+
private final ExecutorService aiExecutor;
2329

2430
// Package-private to allow mock server override in tests
2531
String geminiBaseUrl = "https://generativelanguage.googleapis.com";
@@ -33,8 +39,17 @@ public interface AiCallback {
3339
public AiClient(SettingsManager settings) {
3440
this.settings = settings;
3541
this.httpClient = HttpClient.newBuilder()
36-
.connectTimeout(Duration.ofSeconds(30))
42+
.connectTimeout(Duration.ofSeconds(CONNECT_TIMEOUT_SECONDS))
3743
.build();
44+
this.aiExecutor = Executors.newSingleThreadExecutor(r -> {
45+
Thread t = new Thread(r, "ReconMaster-AI");
46+
t.setDaemon(true);
47+
return t;
48+
});
49+
}
50+
51+
public void shutdown() {
52+
aiExecutor.shutdownNow();
3853
}
3954

4055
/**
@@ -58,15 +73,15 @@ public String sendPrompt(String systemPrompt, String userPrompt) throws Exceptio
5873
* Sends the prompts to the selected AI provider asynchronously.
5974
*/
6075
public void sendPromptAsync(String systemPrompt, String userPrompt, AiCallback callback) {
61-
new Thread(() -> {
76+
aiExecutor.submit(() -> {
6277
try {
6378
String response = sendPrompt(systemPrompt, userPrompt);
6479
callback.onSuccess(response);
6580
} catch (Throwable t) {
6681
logError("Failed to communicate with AI API", t);
6782
callback.onFailure(t);
6883
}
69-
}).start();
84+
});
7085
}
7186

7287
private String sendGoogleGemini(String systemPrompt, String userPrompt) throws Exception {
@@ -209,7 +224,7 @@ private String sendLocalApi(String systemPrompt, String userPrompt) throws Excep
209224
private String sendPostRequest(String url, String requestBody, Map<String, String> headers) throws Exception {
210225
HttpRequest.Builder builder = HttpRequest.newBuilder()
211226
.uri(URI.create(url))
212-
.timeout(Duration.ofSeconds(90))
227+
.timeout(Duration.ofSeconds(REQUEST_TIMEOUT_SECONDS))
213228
.header("Content-Type", "application/json");
214229

215230
if (headers != null) {

‎src/main/java/burp/ui/AiAssistantPanel.java‎

Lines changed: 14 additions & 138 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,15 @@
44
import burp.modules.AiAnalyzer;
55
import burp.modules.AiClient;
66
import burp.ui.utils.CodeBlockPanel;
7+
import burp.ui.utils.MarkdownRenderer;
8+
import burp.ui.utils.ThemeColors;
79
import burp.utils.SettingsManager;
810
import burp.utils.AiProvider;
911

1012
import javax.swing.*;
1113
import javax.swing.border.TitledBorder;
1214
import java.awt.*;
1315
import java.text.SimpleDateFormat;
14-
import java.util.ArrayList;
1516
import java.util.Date;
1617
import java.util.List;
1718
import java.util.function.Supplier;
@@ -166,6 +167,10 @@ public AiAssistantPanel(
166167
add(splitPane, BorderLayout.CENTER);
167168
}
168169

170+
public void shutdown() {
171+
aiClient.shutdown();
172+
}
173+
169174
public void refreshSettings() {
170175
updateOpSecPanel();
171176
}
@@ -374,31 +379,21 @@ private void displayMarkdown(String markdown) {
374379
private void addTextBlock(String md) {
375380
if (md == null || md.trim().isEmpty()) return;
376381

377-
String html = convertMarkdownToHtml(md);
382+
String html = MarkdownRenderer.toHtml(md);
378383

379384
JTextPane textPane = new JTextPane();
380385
textPane.setContentType("text/html");
381386
textPane.setEditable(false);
382387
textPane.setBackground(resultsContainer.getBackground());
383388
textPane.setMargin(new Insets(6, 6, 6, 6));
384389

385-
// Dynamically style based on Burp theme color
386-
Color bg = UIManager.getColor("TextPane.background");
387-
if (bg == null) bg = UIManager.getColor("Panel.background");
388-
if (bg == null) bg = Color.WHITE;
389-
390-
Color fg = UIManager.getColor("TextPane.foreground");
391-
if (fg == null) fg = UIManager.getColor("Label.foreground");
392-
if (fg == null) fg = Color.BLACK;
393-
394-
String bgHex = String.format("#%02x%02x%02x", bg.getRed(), bg.getGreen(), bg.getBlue());
395-
String fgHex = String.format("#%02x%02x%02x", fg.getRed(), fg.getGreen(), fg.getBlue());
396-
397-
boolean isDark = (bg.getRed() + bg.getGreen() + bg.getBlue()) / 3 < 128;
398-
String borderHex = isDark ? "#444444" : "#dddddd";
390+
boolean isDark = ThemeColors.isDark();
391+
String bgHex = ThemeColors.toHex(ThemeColors.background());
392+
String fgHex = ThemeColors.toHex(ThemeColors.foreground());
393+
String borderHex = isDark ? "#444444" : "#dddddd";
399394
String headerBgHex = isDark ? "#2d2d2d" : "#f2f2f2";
400-
String codeBgHex = isDark ? "#282828" : "#f5f5f5";
401-
String evenRowBgHex = isDark ? "#232323" : "#fafafa";
395+
String codeBgHex = isDark ? "#282828" : "#f5f5f5";
396+
String evenRowHex = isDark ? "#232323" : "#fafafa";
402397

403398
String styledHtml = "<html><head><style>" +
404399
"body { font-family: sans-serif; font-size: 11px; color: " + fgHex + "; margin: 0; }" +
@@ -408,7 +403,7 @@ private void addTextBlock(String md) {
408403
"table { border-collapse: collapse; width: 100%; margin-bottom: 8px; margin-top: 4px; }" +
409404
"th { background-color: " + headerBgHex + "; border: 1px solid " + borderHex + "; padding: 4px; text-align: left; font-weight: bold; color: " + fgHex + "; }" +
410405
"td { border: 1px solid " + borderHex + "; padding: 4px; color: " + fgHex + "; }" +
411-
"tr:nth-child(even) { background-color: " + evenRowBgHex + "; }" +
406+
"tr:nth-child(even) { background-color: " + evenRowHex + "; }" +
412407
"ul { margin-top: 2px; margin-bottom: 2px; padding-left: 16px; }" +
413408
"li { margin-bottom: 1px; }" +
414409
"p { margin-top: 2px; margin-bottom: 2px; }" +
@@ -438,125 +433,6 @@ private void addCodeBlock(String language, String code) {
438433
resultsContainer.add(Box.createVerticalStrut(4));
439434
}
440435

441-
private String convertMarkdownToHtml(String md) {
442-
if (md == null) return "";
443-
444-
String[] lines = md.split("\n");
445-
StringBuilder sb = new StringBuilder();
446-
boolean inList = false;
447-
boolean inTable = false;
448-
boolean hasTableHeader = false;
449-
450-
for (String line : lines) {
451-
String trimmed = line.trim();
452-
453-
// List items
454-
if (trimmed.startsWith("- ") || trimmed.startsWith("* ")) {
455-
if (!inList) {
456-
sb.append("<ul>");
457-
inList = true;
458-
}
459-
String content = trimmed.substring(2);
460-
content = replaceMarkdownFormatting(content);
461-
sb.append("<li>").append(content).append("</li>");
462-
continue;
463-
} else {
464-
if (inList) {
465-
sb.append("</ul>");
466-
inList = false;
467-
}
468-
}
469-
470-
// Tables
471-
if (trimmed.startsWith("|") && trimmed.endsWith("|")) {
472-
if (!inTable) {
473-
sb.append("<table>");
474-
inTable = true;
475-
hasTableHeader = false;
476-
}
477-
478-
if (trimmed.matches("^\\|[\\s\\-:\\|]+$")) {
479-
continue; // Skip separator line
480-
}
481-
482-
String[] cells = trimmed.split("\\|");
483-
List<String> validCells = new ArrayList<>();
484-
int startIdx = trimmed.startsWith("|") ? 1 : 0;
485-
int endIdx = cells.length;
486-
if (trimmed.endsWith("|") && cells.length > startIdx) {
487-
endIdx = cells.length - 1;
488-
}
489-
for (int i = startIdx; i < endIdx; i++) {
490-
validCells.add(cells[i].trim());
491-
}
492-
493-
sb.append("<tr>");
494-
for (String cell : validCells) {
495-
String formattedCell = replaceMarkdownFormatting(cell);
496-
if (!hasTableHeader) {
497-
sb.append("<th>").append(formattedCell).append("</th>");
498-
} else {
499-
sb.append("<td>").append(formattedCell).append("</td>");
500-
}
501-
}
502-
sb.append("</tr>");
503-
504-
if (!hasTableHeader) {
505-
hasTableHeader = true;
506-
}
507-
continue;
508-
} else {
509-
if (inTable) {
510-
sb.append("</table>");
511-
inTable = false;
512-
}
513-
}
514-
515-
// Headers
516-
if (trimmed.startsWith("### ")) {
517-
String content = trimmed.substring(4);
518-
sb.append("<h3>").append(replaceMarkdownFormatting(content)).append("</h3>");
519-
continue;
520-
} else if (trimmed.startsWith("## ")) {
521-
String content = trimmed.substring(3);
522-
sb.append("<h2>").append(replaceMarkdownFormatting(content)).append("</h2>");
523-
continue;
524-
} else if (trimmed.startsWith("# ")) {
525-
String content = trimmed.substring(2);
526-
sb.append("<h1>").append(replaceMarkdownFormatting(content)).append("</h1>");
527-
continue;
528-
}
529-
530-
// Regular paragraph
531-
if (!trimmed.isEmpty()) {
532-
sb.append("<p>").append(replaceMarkdownFormatting(trimmed)).append("</p>");
533-
}
534-
}
535-
536-
if (inList) sb.append("</ul>");
537-
if (inTable) sb.append("</table>");
538-
539-
return sb.toString();
540-
}
541-
542-
private String replaceMarkdownFormatting(String text) {
543-
if (text == null) return "";
544-
545-
// Escape HTML
546-
text = text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;");
547-
548-
// Code code
549-
text = text.replaceAll("`([^`]+)`", "<code>$1</code>");
550-
551-
// Bold
552-
text = text.replaceAll("\\*\\*([^\\*]+)\\*\\*", "<b>$1</b>");
553-
554-
// Italic
555-
text = text.replaceAll("\\*([^\\*]+)\\*", "<i>$1</i>");
556-
557-
return text;
558-
}
559-
560436
/**
561437
* Creates a non-editable JEditorPane that renders HTML properly and wraps text
562438
* to the container width (prevents horizontal scrollbar in the left panel).

‎src/main/java/burp/ui/utils/MarkdownRenderer.java‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
package burp.ui.utils;
2+
3+
import java.util.ArrayList;
4+
import java.util.List;
5+
6+
public class MarkdownRenderer {
7+
8+
private MarkdownRenderer() {}
9+
10+
public static String toHtml(String md) {
11+
if (md == null) return "";
12+
13+
String[] lines = md.split("\n");
14+
StringBuilder sb = new StringBuilder();
15+
boolean inList = false;
16+
boolean inTable = false;
17+
boolean hasTableHeader = false;
18+
19+
for (String line : lines) {
20+
String trimmed = line.trim();
21+
22+
// List items
23+
if (trimmed.startsWith("- ") || trimmed.startsWith("* ")) {
24+
if (!inList) {
25+
sb.append("<ul>");
26+
inList = true;
27+
}
28+
sb.append("<li>").append(replaceFormatting(trimmed.substring(2))).append("</li>");
29+
continue;
30+
} else {
31+
if (inList) {
32+
sb.append("</ul>");
33+
inList = false;
34+
}
35+
}
36+
37+
// Tables
38+
if (trimmed.startsWith("|") && trimmed.endsWith("|")) {
39+
if (!inTable) {
40+
sb.append("<table>");
41+
inTable = true;
42+
hasTableHeader = false;
43+
}
44+
if (trimmed.matches("^\\|[\\s\\-:\\|]+$")) continue; // separator row
45+
46+
String[] cells = trimmed.split("\\|");
47+
List<String> validCells = new ArrayList<>();
48+
int start = 1;
49+
int end = cells.length - 1;
50+
for (int i = start; i < end; i++) validCells.add(cells[i].trim());
51+
52+
sb.append("<tr>");
53+
for (String cell : validCells) {
54+
String tag = hasTableHeader ? "td" : "th";
55+
sb.append("<").append(tag).append(">")
56+
.append(replaceFormatting(cell))
57+
.append("</").append(tag).append(">");
58+
}
59+
sb.append("</tr>");
60+
hasTableHeader = true;
61+
continue;
62+
} else {
63+
if (inTable) {
64+
sb.append("</table>");
65+
inTable = false;
66+
}
67+
}
68+
69+
// Headers
70+
if (trimmed.startsWith("### ")) {
71+
sb.append("<h3>").append(replaceFormatting(trimmed.substring(4))).append("</h3>");
72+
} else if (trimmed.startsWith("## ")) {
73+
sb.append("<h2>").append(replaceFormatting(trimmed.substring(3))).append("</h2>");
74+
} else if (trimmed.startsWith("# ")) {
75+
sb.append("<h1>").append(replaceFormatting(trimmed.substring(2))).append("</h1>");
76+
} else if (!trimmed.isEmpty()) {
77+
sb.append("<p>").append(replaceFormatting(trimmed)).append("</p>");
78+
}
79+
}
80+
81+
if (inList) sb.append("</ul>");
82+
if (inTable) sb.append("</table>");
83+
84+
return sb.toString();
85+
}
86+
87+
public static String replaceFormatting(String text) {
88+
if (text == null) return "";
89+
text = text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;");
90+
text = text.replaceAll("`([^`]+)`", "<code>$1</code>");
91+
text = text.replaceAll("\\*\\*([^*]+)\\*\\*", "<b>$1</b>");
92+
text = text.replaceAll("\\*([^*]+)\\*", "<i>$1</i>");
93+
return text;
94+
}
95+
}

0 commit comments

Comments
 (0)