admin/HaE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Version: 4.1.1 Update

Browse Source
This commit is contained in:
gh0stkey
2025-03-25 11:44:07 +08:00
parent 6adf30f25c
commit 124e4c14fd
11 changed files with 167 additions and 306 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/main/java/hae/HaE.java
View File

@@ -2,7 +2,6 @@ package hae;
import burp.api.montoya.BurpExtension;
import burp.api.montoya.MontoyaApi;
import burp.api.montoya.extension.ExtensionUnloadingHandler;
import burp.api.montoya.logging.Logging;
import hae.cache.MessageCache;
import hae.component.Main;
@@ -20,7 +19,7 @@ public class HaE implements BurpExtension {
public void initialize(MontoyaApi api) {
// 设置扩展名称
api.extension().setName("HaE - Highlighter and Extractor");
String version = "4.1";
String version = "4.1.1";
// 加载扩展后输出的项目信息
Logging logging = api.logging();
@@ -53,13 +52,10 @@ public class HaE implements BurpExtension {
dataManager.loadData(messageTableModel);
api.extension().registerUnloadingHandler(new ExtensionUnloadingHandler() {
@Override
public void extensionUnloaded() {
// 卸载清空数据
Config.globalDataMap.clear();
MessageCache.clear();
}
api.extension().registerUnloadingHandler(() -> {
// 卸载清空数据
Config.globalDataMap.clear();
MessageCache.clear();
});
}

46
src/main/java/hae/cache/DataQueryCache.java vendored
View File

@@ -1,46 +0,0 @@
package hae.cache;
import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
public class DataQueryCache {
private static final int MAX_SIZE = 1000;
private static final int EXPIRE_DURATION = 30;
private static final Cache<String, Map<String, List<String>>> hostQueryCache =
Caffeine.newBuilder()
.maximumSize(MAX_SIZE)
.expireAfterWrite(EXPIRE_DURATION, TimeUnit.MINUTES)
.build();
private static final Cache<String, List<String>> hostFilterCache =
Caffeine.newBuilder()
.maximumSize(MAX_SIZE)
.expireAfterWrite(EXPIRE_DURATION, TimeUnit.MINUTES)
.build();
public static void putHostQueryResult(String host, Map<String, List<String>> result) {
hostQueryCache.put(host, result);
}
public static Map<String, List<String>> getHostQueryResult(String host) {
return hostQueryCache.getIfPresent(host);
}
public static void putHostFilterResult(String input, List<String> result) {
hostFilterCache.put(input, result);
}
public static List<String> getHostFilterResult(String input) {
return hostFilterCache.getIfPresent(input);
}
public static void clearCache() {
hostQueryCache.invalidateAll();
hostFilterCache.invalidateAll();
}
}

10
src/main/java/hae/cache/MessageCache.java vendored
View File

@@ -4,16 +4,10 @@ import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import java.util.Map;
import java.util.concurrent.TimeUnit;
public class MessageCache {
private static final int MAX_SIZE = 100000;
private static final int EXPIRE_DURATION = 5;
private static final Cache<String, Map<String, Map<String, Object>>> cache =
Caffeine.newBuilder()
.maximumSize(MAX_SIZE)
.expireAfterWrite(EXPIRE_DURATION, TimeUnit.HOURS)
.build();
public static void put(String key, Map<String, Map<String, Object>> value) {
@@ -24,10 +18,6 @@ public class MessageCache {
return cache.getIfPresent(key);
}
public static void remove(String key) {
cache.invalidate(key);
}
public static void clear() {
cache.invalidateAll();
}

28
src/main/java/hae/component/Config.java
View File

@@ -14,7 +14,6 @@ import javax.swing.border.EmptyBorder;
import javax.swing.border.TitledBorder;
import javax.swing.event.DocumentEvent;
import javax.swing.event.DocumentListener;
import javax.swing.event.TableModelEvent;
import javax.swing.event.TableModelListener;
import javax.swing.table.DefaultTableModel;
import java.awt.*;
@@ -174,12 +173,10 @@ public class Config extends JPanel {
}
private TableModelListener craeteSettingTableModelListener(JComboBox<String> setTypeComboBox, DefaultTableModel model) {
return new TableModelListener() {
@Override
public void tableChanged(TableModelEvent e) {
String selected = (String) setTypeComboBox.getSelectedItem();
String values = getFirstColumnDataAsString(model);
return e -> {
String selected = (String) setTypeComboBox.getSelectedItem();
String values = getFirstColumnDataAsString(model);
if (selected != null) {
if (selected.equals("Exclude suffix")) {
if (!values.equals(configLoader.getExcludeSuffix()) && !values.isEmpty()) {
configLoader.setExcludeSuffix(values);
@@ -197,18 +194,15 @@ public class Config extends JPanel {
configLoader.setExcludeStatus(values);
}
}
}
};
}
private ActionListener createSettingActionListener(JComboBox<String> setTypeComboBox, DefaultTableModel model) {
return new ActionListener() {
@Override
public void actionPerformed(ActionEvent e) {
String selected = (String) setTypeComboBox.getSelectedItem();
model.setRowCount(0);
return e -> {
String selected = (String) setTypeComboBox.getSelectedItem();
model.setRowCount(0);
if (selected != null) {
if (selected.equals("Exclude suffix")) {
addDataToTable(configLoader.getExcludeSuffix().replaceAll("\\|", "\r\n"), model);
}
@@ -286,13 +280,13 @@ public class Config extends JPanel {
settingPanel.add(inputPanel, BorderLayout.CENTER);
addButton.addActionListener(e -> addActionPerformed(e, model, addTextField, setTypeComboBox.getSelectedItem().toString()));
addButton.addActionListener(e -> addActionPerformed(e, model, addTextField));
addTextField.addKeyListener(new KeyAdapter() {
@Override
public void keyPressed(KeyEvent e) {
if (e.getKeyCode() == KeyEvent.VK_ENTER) {
addActionPerformed(null, model, addTextField, setTypeComboBox.getSelectedItem().toString());
addActionPerformed(null, model, addTextField);
}
}
});
@@ -413,7 +407,7 @@ public class Config extends JPanel {
configLoader.setScope(String.join("|", HaEScope));
}
private void addActionPerformed(ActionEvent e, DefaultTableModel model, JTextField addTextField, String comboBoxSelected) {
private void addActionPerformed(ActionEvent e, DefaultTableModel model, JTextField addTextField) {
String addTextFieldText = addTextField.getText();
if (addTextField.getForeground().equals(Color.BLACK)) {
addDataToTable(addTextFieldText, model);

55
src/main/java/hae/component/board/Databoard.java
View File

@@ -2,7 +2,6 @@ package hae.component.board;
import burp.api.montoya.MontoyaApi;
import hae.Config;
import hae.cache.DataQueryCache;
import hae.component.board.message.MessageTableModel;
import hae.component.board.message.MessageTableModel.MessageTable;
import hae.component.board.table.Datatable;
@@ -46,19 +45,6 @@ public class Databoard extends JPanel {
initComponents();
}
public static void setProgressBar(boolean status, JProgressBar progressBar, String showString) {
progressBar.setIndeterminate(status);
if (!status) {
progressBar.setMaximum(100);
progressBar.setString("OK");
progressBar.setStringPainted(true);
progressBar.setValue(progressBar.getMaximum());
} else {
progressBar.setString(showString);
progressBar.setStringPainted(true);
}
}
private void initComponents() {
setLayout(new GridBagLayout());
((GridBagLayout) getLayout()).columnWidths = new int[]{25, 0, 0, 0, 20, 0};
@@ -136,7 +122,16 @@ public class Databoard extends JPanel {
}
private void setProgressBar(boolean status) {
setProgressBar(status, progressBar, "Loading ...");
progressBar.setIndeterminate(status);
if (!status) {
progressBar.setMaximum(100);
progressBar.setString("OK");
progressBar.setStringPainted(true);
progressBar.setValue(progressBar.getMaximum());
} else {
progressBar.setString("Loading...");
progressBar.setStringPainted(true);
}
}
private void setAutoMatch() {
@@ -182,7 +177,7 @@ public class Databoard extends JPanel {
handleComboBoxWorker.cancel(true);
}
handleComboBoxWorker = new SwingWorker<Map<String, List<String>>, Void>() {
handleComboBoxWorker = new SwingWorker<>() {
@Override
protected Map<String, List<String>> doInBackground() {
return getSelectedMapByHost(selectedHost);
@@ -253,12 +248,6 @@ public class Databoard extends JPanel {
}
private Map<String, List<String>> getSelectedMapByHost(String selectedHost) {
// 先尝试从缓存获取结果
Map<String, List<String>> cachedResult = DataQueryCache.getHostQueryResult(selectedHost);
if (cachedResult != null) {
return cachedResult;
}
ConcurrentHashMap<String, Map<String, List<String>>> dataMap = Config.globalDataMap;
Map<String, List<String>> selectedDataMap;
@@ -284,11 +273,6 @@ public class Databoard extends JPanel {
selectedDataMap = dataMap.get(selectedHost);
}
// 将结果存入缓存
if (selectedDataMap != null) {
DataQueryCache.putHostQueryResult(selectedHost, selectedDataMap);
}
return selectedDataMap;
}
@@ -323,10 +307,10 @@ public class Databoard extends JPanel {
applyHostFilterWorker.cancel(true);
}
applyHostFilterWorker = new SwingWorker<Void, Void>() {
applyHostFilterWorker = new SwingWorker<>() {
@Override
protected Void doInBackground() throws Exception {
RowFilter<Object, Object> rowFilter = new RowFilter<Object, Object>() {
protected Void doInBackground() {
RowFilter<Object, Object> rowFilter = new RowFilter<>() {
public boolean include(Entry<?, ?> entry) {
if (cleanedText.equals("*")) {
return true;
@@ -348,24 +332,15 @@ public class Databoard extends JPanel {
}
private List<String> getHostByList() {
// 先尝试从缓存获取结果
List<String> cachedResult = DataQueryCache.getHostFilterResult("all_hosts");
if (cachedResult != null) {
return cachedResult;
}
List<String> result = new ArrayList<>();
if (!Config.globalDataMap.isEmpty()) {
result = new ArrayList<>(Config.globalDataMap.keySet());
// 将结果存入缓存
DataQueryCache.putHostFilterResult("all_hosts", result);
}
return result;
}
private void clearActionPerformed(ActionEvent e) {
// 清除缓存
DataQueryCache.clearCache();
int retCode = JOptionPane.showConfirmDialog(this, "Do you want to clear data?", "Info",
JOptionPane.YES_NO_OPTION);
String host = hostTextField.getText();

231
src/main/java/hae/component/board/message/MessageTableModel.java
View File

@@ -10,10 +10,8 @@ import burp.api.montoya.ui.UserInterface;
import burp.api.montoya.ui.editor.HttpRequestEditor;
import burp.api.montoya.ui.editor.HttpResponseEditor;
import hae.Config;
import hae.cache.MessageCache;
import hae.utils.ConfigLoader;
import hae.utils.DataManager;
import hae.utils.string.HashCalculator;
import hae.utils.string.StringProcessor;
import javax.swing.*;
@@ -58,14 +56,25 @@ public class MessageTableModel extends AbstractTableModel {
messageTable.setAutoCreateRowSorter(true);
// Length字段根据大小进行排序
TableRowSorter<DefaultTableModel> sorter = getDefaultTableModelTableRowSorter();
messageTable.setRowSorter(sorter);
messageTable.setAutoResizeMode(JTable.AUTO_RESIZE_OFF);
splitPane = new JSplitPane(JSplitPane.VERTICAL_SPLIT);
// 请求/响应文本框
JScrollPane scrollPane = new JScrollPane(messageTable);
scrollPane.setHorizontalScrollBarPolicy(JScrollPane.HORIZONTAL_SCROLLBAR_ALWAYS);
scrollPane.setVerticalScrollBarPolicy(JScrollPane.VERTICAL_SCROLLBAR_ALWAYS);
splitPane.setLeftComponent(scrollPane);
splitPane.setRightComponent(messageTab);
}
private TableRowSorter<DefaultTableModel> getDefaultTableModelTableRowSorter() {
TableRowSorter<DefaultTableModel> sorter = (TableRowSorter<DefaultTableModel>) messageTable.getRowSorter();
sorter.setComparator(4, new Comparator<String>() {
@Override
public int compare(String s1, String s2) {
Integer age1 = Integer.parseInt(s1);
Integer age2 = Integer.parseInt(s2);
return age1.compareTo(age2);
}
sorter.setComparator(4, (Comparator<String>) (s1, s2) -> {
Integer age1 = Integer.parseInt(s1);
Integer age2 = Integer.parseInt(s2);
return age1.compareTo(age2);
});
// Color字段根据颜色顺序进行排序
@@ -86,48 +95,31 @@ public class MessageTableModel extends AbstractTableModel {
return -1;
}
});
messageTable.setRowSorter(sorter);
messageTable.setAutoResizeMode(JTable.AUTO_RESIZE_OFF);
splitPane = new JSplitPane(JSplitPane.VERTICAL_SPLIT);
// 请求/响应文本框
JScrollPane scrollPane = new JScrollPane(messageTable);
scrollPane.setHorizontalScrollBarPolicy(JScrollPane.HORIZONTAL_SCROLLBAR_ALWAYS);
scrollPane.setVerticalScrollBarPolicy(JScrollPane.VERTICAL_SCROLLBAR_ALWAYS);
splitPane.setLeftComponent(scrollPane);
splitPane.setRightComponent(messageTab);
return sorter;
}
public synchronized void add(HttpRequestResponse messageInfo, String url, String method, String status, String length, String comment, String color, boolean flag) {
synchronized (log) {
boolean isDuplicate = false;
MessageEntry logEntry = new MessageEntry(messageInfo, method, url, comment, length, color, status);
byte[] reqByteA = new byte[0];
byte[] resByteA = new byte[0];
if (messageInfo != null) {
HttpRequest httpRequest = messageInfo.request();
HttpResponse httpResponse = messageInfo.response();
reqByteA = httpRequest.toByteArray().getBytes();
resByteA = httpResponse.toByteArray().getBytes();
if (messageInfo == null) {
return;
}
// 比较Hash，如若存在重复的请求或响应，则不放入消息内容里
boolean isDuplicate = false;
try {
if (!log.isEmpty()) {
if (!log.isEmpty() && flag) {
String host = StringProcessor.getHostByUrl(url);
for (MessageEntry entry : log) {
HttpRequestResponse reqResMessage = entry.getRequestResponse();
byte[] reqByteB = reqResMessage.request().toByteArray().getBytes();
byte[] resByteB = reqResMessage.response().toByteArray().getBytes();
try {
// 通过URL、请求和响应报文、匹配数据内容，多维度进行对比
if ((entry.getUrl().equals(url) || (Arrays.equals(reqByteB, reqByteA) || Arrays.equals(resByteB, resByteA))) && (areMapsEqual(getCacheData(reqByteB), getCacheData(reqByteA)) && areMapsEqual(getCacheData(resByteB), getCacheData(resByteA)))) {
if (host.equals(StringProcessor.getHostByUrl(entry.getUrl()))) {
if (isRequestDuplicate(
messageInfo, entry.getRequestResponse(),
url, entry.getUrl(),
comment, entry.getComment(),
color, entry.getColor()
)) {
isDuplicate = true;
break;
}
} catch (Exception ignored) {
}
}
}
@@ -136,42 +128,82 @@ public class MessageTableModel extends AbstractTableModel {
if (!isDuplicate) {
if (flag) {
try {
DataManager dataManager = new DataManager(api);
// 数据存储在BurpSuite空间内
PersistedObject persistedObject = PersistedObject.persistedObject();
persistedObject.setHttpRequestResponse("messageInfo", messageInfo);
persistedObject.setString("comment", comment);
persistedObject.setString("color", color);
String uuidIndex = StringProcessor.getRandomUUID();
dataManager.putData("message", uuidIndex, persistedObject);
} catch (Exception ignored) {
}
persistData(messageInfo, comment, color);
}
// 添加进日志
log.add(logEntry);
log.add(new MessageEntry(messageInfo, method, url, comment, length, color, status));
}
}
}
public synchronized void addBatch(List<Object[]> batchData) {
synchronized (log) {
for (Object[] data : batchData) {
HttpRequestResponse messageInfo = (HttpRequestResponse) data[0];
String url = (String) data[1];
String method = (String) data[2];
String status = (String) data[3];
String length = (String) data[4];
String comment = (String) data[5];
String color = (String) data[6];
private boolean isRequestDuplicate(
HttpRequestResponse newReq, HttpRequestResponse existingReq,
String newUrl, String existingUrl,
String newComment, String existingComment,
String newColor, String existingColor) {
try {
// 基础属性匹配
String normalizedNewUrl = normalizeUrl(newUrl);
String normalizedExistingUrl = normalizeUrl(existingUrl);
boolean basicMatch = normalizedNewUrl.equals(normalizedExistingUrl);
// 复用现有的 add 方法逻辑，但跳过重复检查
MessageEntry logEntry = new MessageEntry(messageInfo, method, url, comment, length, color, status);
log.add(logEntry);
}
// 请求响应内容匹配
byte[] newReqBytes = newReq.request().toByteArray().getBytes();
byte[] newResBytes = newReq.response().toByteArray().getBytes();
byte[] existingReqBytes = existingReq.request().toByteArray().getBytes();
byte[] existingResBytes = existingReq.response().toByteArray().getBytes();
boolean contentMatch = Arrays.equals(newReqBytes, existingReqBytes) &&
Arrays.equals(newResBytes, existingResBytes);
// 注释和颜色匹配
boolean metadataMatch = areCommentsEqual(newComment, existingComment) &&
newColor.equals(existingColor);
return (basicMatch || contentMatch) && metadataMatch;
} catch (Exception e) {
return false;
}
}
private String normalizeUrl(String url) {
if (url == null) {
return "";
}
String normalized = url.trim().toLowerCase();
while (normalized.endsWith("/")) {
normalized = normalized.substring(0, normalized.length() - 1);
}
return normalized.replaceAll("//", "/");
}
private boolean areCommentsEqual(String comment1, String comment2) {
if (comment1 == null || comment2 == null) {
return false;
}
try {
// 将注释按规则拆分并排序
Set<String> rules1 = new TreeSet<>(Arrays.asList(comment1.split(", ")));
Set<String> rules2 = new TreeSet<>(Arrays.asList(comment2.split(", ")));
return rules1.equals(rules2);
} catch (Exception e) {
return false;
}
}
private void persistData(HttpRequestResponse messageInfo, String comment, String color) {
try {
DataManager dataManager = new DataManager(api);
PersistedObject persistedObject = PersistedObject.persistedObject();
persistedObject.setHttpRequestResponse("messageInfo", messageInfo);
persistedObject.setString("comment", comment);
persistedObject.setString("color", color);
String uuidIndex = StringProcessor.getRandomUUID();
dataManager.putData("message", uuidIndex, persistedObject);
} catch (Exception e) {
api.logging().logToError("Data persistence error: " + e.getMessage());
}
}
@@ -183,7 +215,7 @@ public class MessageTableModel extends AbstractTableModel {
currentWorker.cancel(true);
}
currentWorker = new SwingWorker<Void, Void>() {
currentWorker = new SwingWorker<>() {
@Override
protected Void doInBackground() {
for (int i = 0; i < log.size(); i++) {
@@ -333,56 +365,6 @@ public class MessageTableModel extends AbstractTableModel {
return isMatch;
}
private Map<String, Map<String, Object>> getCacheData(byte[] content) {
String hashIndex = HashCalculator.calculateHash(content);
return MessageCache.get(hashIndex);
}
private boolean areMapsEqual(Map<String, Map<String, Object>> map1, Map<String, Map<String, Object>> map2) {
if (map1 == null || map2 == null) {
return false;
}
if (map1.size() != map2.size()) {
return false;
}
for (String key : map1.keySet()) {
if (!map2.containsKey(key)) {
return false;
}
if (areInnerMapsEqual(map1.get(key), map2.get(key))) {
return false;
}
}
return true;
}
private boolean areInnerMapsEqual(Map<String, Object> innerMap1, Map<String, Object> innerMap2) {
if (innerMap1.size() != innerMap2.size()) {
return true;
}
for (String key : innerMap1.keySet()) {
if (!innerMap2.containsKey(key)) {
return true;
}
Object value1 = innerMap1.get(key);
Object value2 = innerMap2.get(key);
// 如果值是Map，则递归对比
if (value1 instanceof Map && value2 instanceof Map) {
if (areInnerMapsEqual((Map<String, Object>) value1, (Map<String, Object>) value2)) {
return true;
}
} else if (!value1.equals(value2)) {
return true;
}
}
return false;
}
public JSplitPane getSplitPane() {
return splitPane;
}
@@ -391,10 +373,6 @@ public class MessageTableModel extends AbstractTableModel {
return messageTable;
}
public LinkedList<MessageEntry> getLogs() {
return log;
}
@Override
public int getRowCount() {
return filteredLog.size();
@@ -447,7 +425,6 @@ public class MessageTableModel extends AbstractTableModel {
private final ExecutorService executorService;
private final HttpRequestEditor requestEditor;
private final HttpResponseEditor responseEditor;
private MessageEntry messageEntry;
private int lastSelectedIndex = -1;
public MessageTable(TableModel messageTableModel, HttpRequestEditor requestEditor, HttpResponseEditor responseEditor) {
@@ -468,7 +445,7 @@ public class MessageTableModel extends AbstractTableModel {
}
private void getSelectedMessage() {
messageEntry = filteredLog.get(lastSelectedIndex);
MessageEntry messageEntry = filteredLog.get(lastSelectedIndex);
HttpRequestResponse httpRequestResponse = messageEntry.getRequestResponse();

9
src/main/java/hae/component/board/table/Datatable.java
View File

@@ -55,12 +55,7 @@ public class Datatable extends JPanel {
dataTable.setRowSorter(sorter);
// 设置ID排序
sorter.setComparator(0, new Comparator<Integer>() {
@Override
public int compare(Integer s1, Integer s2) {
return s1.compareTo(s2);
}
});
sorter.setComparator(0, (Comparator<Integer>) Integer::compareTo);
for (String item : dataList) {
if (!item.isEmpty()) {
@@ -180,7 +175,7 @@ public class Datatable extends JPanel {
}
private RowFilter<Object, Object> getObjectObjectRowFilter(JTextField searchField, boolean firstFlag) {
return new RowFilter<Object, Object>() {
return new RowFilter<>() {
public boolean include(Entry<?, ?> entry) {
String searchFieldTextText = searchField.getText();
searchFieldTextText = searchFieldTextText.toLowerCase();

4
src/main/java/hae/component/rule/Rule.java
View File

@@ -97,7 +97,7 @@ public class Rule extends JPanel {
Display ruleDisplay = new Display();
ruleDisplay.formatTextField.setText("{0}");
int showState = JOptionPane.showConfirmDialog(this, ruleDisplay, "Add Rule", JOptionPane.OK_OPTION);
int showState = JOptionPane.showConfirmDialog(this, ruleDisplay, "Add Rule", JOptionPane.YES_NO_OPTION);
if (showState == YES_OPTION) {
Vector<Object> ruleData = new Vector<>();
ruleData.add(false);
@@ -132,7 +132,7 @@ public class Rule extends JPanel {
ruleDisplay.formatTextField.setEnabled(ruleDisplay.engineComboBox.getSelectedItem().toString().equals("nfa"));
int showState = JOptionPane.showConfirmDialog(this, ruleDisplay, "Edit Rule", JOptionPane.OK_OPTION);
int showState = JOptionPane.showConfirmDialog(this, ruleDisplay, "Edit Rule", JOptionPane.YES_NO_OPTION);
if (showState == 0) {
int select = ruleTable.convertRowIndexToModel(ruleTable.getSelectedRow());
model.setValueAt(ruleDisplay.ruleNameTextField.getText(), select, 1);

2
src/main/java/hae/component/rule/Rules.java
View File

@@ -59,8 +59,6 @@ public class Rules extends JTabbedPane {
private void initComponents() {
reloadRuleGroup();
JTabbedPane tabbedPane = this;
JMenuItem deleteMenuItem = new JMenuItem("Delete");
JPopupMenu popupMenu = new JPopupMenu();
popupMenu.add(deleteMenuItem);

4
src/main/java/hae/instances/http/utils/MessageProcessor.java
View File

@@ -77,10 +77,10 @@ public class MessageProcessor {
List<String> commentList = resultList.get(1);
if (!colorList.isEmpty() && !commentList.isEmpty()) {
String color = retrieveFinalColor(retrieveColorIndices(colorList));
Map<String, String> colorMap = new HashMap<String, String>() {{
Map<String, String> colorMap = new HashMap<>() {{
put("color", color);
}};
Map<String, String> commentMap = new HashMap<String, String>() {{
Map<String, String> commentMap = new HashMap<>() {{
put("comment", String.join(", ", commentList));
}};
highlightList.add(colorMap);

70
src/main/java/hae/utils/DataManager.java
View File

@@ -10,11 +10,10 @@ import burp.api.montoya.persistence.Persistence;
import hae.component.board.message.MessageTableModel;
import hae.instances.http.utils.RegularMatcher;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
public class DataManager {
private final MontoyaApi api;
@@ -65,9 +64,7 @@ public class DataManager {
dataIndex.forEach(index -> {
PersistedObject dataObj = persistence.extensionData().getChildObject(index);
try {
dataObj.stringListKeys().forEach(dataKey -> {
RegularMatcher.putDataToGlobalMap(api, index, dataKey, dataObj.getStringList(dataKey).stream().toList(), false);
});
dataObj.stringListKeys().forEach(dataKey -> RegularMatcher.putDataToGlobalMap(api, index, dataKey, dataObj.getStringList(dataKey).stream().toList(), false));
} catch (Exception ignored) {
}
});
@@ -79,69 +76,54 @@ public class DataManager {
return;
}
List<String> indexList = new ArrayList<>();
for (Object item : messageIndex) {
try {
if (item != null) {
indexList.add(item.toString());
}
} catch (Exception e) {
api.logging().logToError("转换索引时出错: " + e.getMessage());
}
// 直接转换为List，简化处理
List<String> indexList = messageIndex.stream()
.filter(Objects::nonNull)
.map(Object::toString)
.toList();
if (indexList.isEmpty()) {
return;
}
final int batchSize = 2000; // 增加批处理大小
final int threadCount = Math.max(8, Runtime.getRuntime().availableProcessors() * 2); // 增加线程数
int totalSize = indexList.size();
// 使用更高效的线程池
final int batchSize = 2000;
final int threadCount = Math.max(8, Runtime.getRuntime().availableProcessors() * 2);
ExecutorService executorService = Executors.newWorkStealingPool(threadCount);
List<Future<List<Object[]>>> futures = new ArrayList<>();
// 分批并行处理数据
for (int i = 0; i < totalSize; i += batchSize) {
int endIndex = Math.min(i + batchSize, totalSize);
List<String> batch = indexList.subList(i, endIndex);
Future<List<Object[]>> future = executorService.submit(() -> processBatchParallel(batch));
futures.add(future);
}
// 批量添加数据到模型
try {
for (Future<List<Object[]>> future : futures) {
List<Object[]> batchData = future.get();
messageTableModel.addBatch(batchData);
// 分批处理
for (int i = 0; i < indexList.size(); i += batchSize) {
int endIndex = Math.min(i + batchSize, indexList.size());
List<String> batch = indexList.subList(i, endIndex);
processBatch(batch, messageTableModel);
}
} catch (Exception e) {
api.logging().logToError("批量添加数据时出错: " + e.getMessage());
} finally {
executorService.shutdown();
}
}
private List<Object[]> processBatchParallel(List<String> batch) {
List<Object[]> batchData = new ArrayList<>();
for (String index : batch) {
private void processBatch(List<String> batch, MessageTableModel messageTableModel) {
batch.forEach(index -> {
try {
PersistedObject dataObj = persistence.extensionData().getChildObject(index);
if (dataObj != null) {
HttpRequestResponse messageInfo = dataObj.getHttpRequestResponse("messageInfo");
if (messageInfo != null) {
batchData.add(prepareMessageData(messageInfo, dataObj));
addMessageToModel(messageInfo, dataObj, messageTableModel);
}
}
} catch (Exception e) {
api.logging().logToError("处理消息数据时出错: " + e.getMessage() + ", index: " + index);
}
}
return batchData;
});
}
private Object[] prepareMessageData(HttpRequestResponse messageInfo, PersistedObject dataObj) {
private void addMessageToModel(HttpRequestResponse messageInfo, PersistedObject dataObj, MessageTableModel messageTableModel) {
HttpRequest request = messageInfo.request();
HttpResponse response = messageInfo.response();
return new Object[]{
messageTableModel.add(
messageInfo,
request.url(),
request.method(),
@@ -150,6 +132,6 @@ public class DataManager {
dataObj.getString("comment"),
dataObj.getString("color"),
false
};
);
}
}