diff --git a/src/main/java/hae/HaE.java b/src/main/java/hae/HaE.java index 5a6cee9..4ad96d4 100644 --- a/src/main/java/hae/HaE.java +++ b/src/main/java/hae/HaE.java @@ -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(); }); } diff --git a/src/main/java/hae/cache/DataQueryCache.java b/src/main/java/hae/cache/DataQueryCache.java deleted file mode 100644 index 85719eb..0000000 --- a/src/main/java/hae/cache/DataQueryCache.java +++ /dev/null @@ -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>> hostQueryCache = - Caffeine.newBuilder() - .maximumSize(MAX_SIZE) - .expireAfterWrite(EXPIRE_DURATION, TimeUnit.MINUTES) - .build(); - - private static final Cache> hostFilterCache = - Caffeine.newBuilder() - .maximumSize(MAX_SIZE) - .expireAfterWrite(EXPIRE_DURATION, TimeUnit.MINUTES) - .build(); - - public static void putHostQueryResult(String host, Map> result) { - hostQueryCache.put(host, result); - } - - public static Map> getHostQueryResult(String host) { - return hostQueryCache.getIfPresent(host); - } - - public static void putHostFilterResult(String input, List result) { - hostFilterCache.put(input, result); - } - - public static List getHostFilterResult(String input) { - return hostFilterCache.getIfPresent(input); - } - - public static void clearCache() { - hostQueryCache.invalidateAll(); - hostFilterCache.invalidateAll(); - } -} \ No newline at end of file diff --git a/src/main/java/hae/cache/MessageCache.java b/src/main/java/hae/cache/MessageCache.java index 55973be..0efa7db 100644 --- a/src/main/java/hae/cache/MessageCache.java +++ b/src/main/java/hae/cache/MessageCache.java @@ -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>> cache = Caffeine.newBuilder() - .maximumSize(MAX_SIZE) - .expireAfterWrite(EXPIRE_DURATION, TimeUnit.HOURS) .build(); public static void put(String key, Map> 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(); } diff --git a/src/main/java/hae/component/Config.java b/src/main/java/hae/component/Config.java index fd802bc..8f21f0b 100644 --- a/src/main/java/hae/component/Config.java +++ b/src/main/java/hae/component/Config.java @@ -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 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 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); diff --git a/src/main/java/hae/component/board/Databoard.java b/src/main/java/hae/component/board/Databoard.java index 043a21b..e1fc0a0 100644 --- a/src/main/java/hae/component/board/Databoard.java +++ b/src/main/java/hae/component/board/Databoard.java @@ -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>, Void>() { + handleComboBoxWorker = new SwingWorker<>() { @Override protected Map> doInBackground() { return getSelectedMapByHost(selectedHost); @@ -253,12 +248,6 @@ public class Databoard extends JPanel { } private Map> getSelectedMapByHost(String selectedHost) { - // 先尝试从缓存获取结果 - Map> cachedResult = DataQueryCache.getHostQueryResult(selectedHost); - if (cachedResult != null) { - return cachedResult; - } - ConcurrentHashMap>> dataMap = Config.globalDataMap; Map> 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() { + applyHostFilterWorker = new SwingWorker<>() { @Override - protected Void doInBackground() throws Exception { - RowFilter rowFilter = new RowFilter() { + protected Void doInBackground() { + RowFilter rowFilter = new RowFilter<>() { public boolean include(Entry entry) { if (cleanedText.equals("*")) { return true; @@ -348,24 +332,15 @@ public class Databoard extends JPanel { } private List getHostByList() { - // 先尝试从缓存获取结果 - List cachedResult = DataQueryCache.getHostFilterResult("all_hosts"); - if (cachedResult != null) { - return cachedResult; - } - List 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(); diff --git a/src/main/java/hae/component/board/message/MessageTableModel.java b/src/main/java/hae/component/board/message/MessageTableModel.java index 409dc14..a2eb12e 100644 --- a/src/main/java/hae/component/board/message/MessageTableModel.java +++ b/src/main/java/hae/component/board/message/MessageTableModel.java @@ -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 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 getDefaultTableModelTableRowSorter() { TableRowSorter sorter = (TableRowSorter) messageTable.getRowSorter(); - sorter.setComparator(4, new Comparator() { - @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) (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 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 rules1 = new TreeSet<>(Arrays.asList(comment1.split(", "))); + Set 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() { + 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> getCacheData(byte[] content) { - String hashIndex = HashCalculator.calculateHash(content); - return MessageCache.get(hashIndex); - } - - private boolean areMapsEqual(Map> map1, Map> 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 innerMap1, Map 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) value1, (Map) 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 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(); diff --git a/src/main/java/hae/component/board/table/Datatable.java b/src/main/java/hae/component/board/table/Datatable.java index 207c9ee..e08e61a 100644 --- a/src/main/java/hae/component/board/table/Datatable.java +++ b/src/main/java/hae/component/board/table/Datatable.java @@ -55,12 +55,7 @@ public class Datatable extends JPanel { dataTable.setRowSorter(sorter); // 设置ID排序 - sorter.setComparator(0, new Comparator() { - @Override - public int compare(Integer s1, Integer s2) { - return s1.compareTo(s2); - } - }); + sorter.setComparator(0, (Comparator) Integer::compareTo); for (String item : dataList) { if (!item.isEmpty()) { @@ -180,7 +175,7 @@ public class Datatable extends JPanel { } private RowFilter getObjectObjectRowFilter(JTextField searchField, boolean firstFlag) { - return new RowFilter() { + return new RowFilter<>() { public boolean include(Entry entry) { String searchFieldTextText = searchField.getText(); searchFieldTextText = searchFieldTextText.toLowerCase(); diff --git a/src/main/java/hae/component/rule/Rule.java b/src/main/java/hae/component/rule/Rule.java index 49078a7..d4ed8d2 100644 --- a/src/main/java/hae/component/rule/Rule.java +++ b/src/main/java/hae/component/rule/Rule.java @@ -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 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); diff --git a/src/main/java/hae/component/rule/Rules.java b/src/main/java/hae/component/rule/Rules.java index f2765ca..e5c5026 100644 --- a/src/main/java/hae/component/rule/Rules.java +++ b/src/main/java/hae/component/rule/Rules.java @@ -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); diff --git a/src/main/java/hae/instances/http/utils/MessageProcessor.java b/src/main/java/hae/instances/http/utils/MessageProcessor.java index 16ba4c7..d4c4cb6 100644 --- a/src/main/java/hae/instances/http/utils/MessageProcessor.java +++ b/src/main/java/hae/instances/http/utils/MessageProcessor.java @@ -77,10 +77,10 @@ public class MessageProcessor { List commentList = resultList.get(1); if (!colorList.isEmpty() && !commentList.isEmpty()) { String color = retrieveFinalColor(retrieveColorIndices(colorList)); - Map colorMap = new HashMap() {{ + Map colorMap = new HashMap<>() {{ put("color", color); }}; - Map commentMap = new HashMap() {{ + Map commentMap = new HashMap<>() {{ put("comment", String.join(", ", commentList)); }}; highlightList.add(colorMap); diff --git a/src/main/java/hae/utils/DataManager.java b/src/main/java/hae/utils/DataManager.java index 0749ab2..9b30240 100644 --- a/src/main/java/hae/utils/DataManager.java +++ b/src/main/java/hae/utils/DataManager.java @@ -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 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 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>> futures = new ArrayList<>(); - // 分批并行处理数据 - for (int i = 0; i < totalSize; i += batchSize) { - int endIndex = Math.min(i + batchSize, totalSize); - List batch = indexList.subList(i, endIndex); - - Future> future = executorService.submit(() -> processBatchParallel(batch)); - futures.add(future); - } - - // 批量添加数据到模型 try { - for (Future> future : futures) { - List batchData = future.get(); - messageTableModel.addBatch(batchData); + // 分批处理 + for (int i = 0; i < indexList.size(); i += batchSize) { + int endIndex = Math.min(i + batchSize, indexList.size()); + List batch = indexList.subList(i, endIndex); + + processBatch(batch, messageTableModel); } - } catch (Exception e) { - api.logging().logToError("批量添加数据时出错: " + e.getMessage()); } finally { executorService.shutdown(); } } - private List processBatchParallel(List batch) { - List batchData = new ArrayList<>(); - for (String index : batch) { + private void processBatch(List 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 - }; + ); } } \ No newline at end of file