From 20afa30822f1397d2351f898dc492a3a2ced8a1c Mon Sep 17 00:00:00 2001 From: gh0stkey <24655118+gh0stkey@users.noreply.github.com> Date: Fri, 21 Mar 2025 21:22:11 +0800 Subject: [PATCH] Version: 4.1 Update --- src/main/java/hae/HaE.java | 6 +- src/main/java/hae/cache/DataQueryCache.java | 46 ++++++ .../{CachePool.java => MessageCache.java} | 2 +- .../java/hae/component/board/Databoard.java | 78 ++++++++-- .../board/message/MessageTableModel.java | 32 +++- src/main/java/hae/component/rule/Rules.java | 138 +++++++++--------- .../hae/instances/editor/RequestEditor.java | 47 +++--- .../hae/instances/editor/ResponseEditor.java | 3 +- .../hae/instances/editor/WebSocketEditor.java | 3 +- .../http/utils/MessageProcessor.java | 1 - .../instances/http/utils/RegularMatcher.java | 100 ++++++------- src/main/java/hae/utils/ConfigLoader.java | 66 ++++----- src/main/java/hae/utils/DataManager.java | 91 ++++++++++-- 13 files changed, 393 insertions(+), 220 deletions(-) create mode 100644 src/main/java/hae/cache/DataQueryCache.java rename src/main/java/hae/cache/{CachePool.java => MessageCache.java} (97%) diff --git a/src/main/java/hae/HaE.java b/src/main/java/hae/HaE.java index ad7e90c..6215cd6 100644 --- a/src/main/java/hae/HaE.java +++ b/src/main/java/hae/HaE.java @@ -4,7 +4,7 @@ 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.CachePool; +import hae.cache.MessageCache; import hae.component.Main; import hae.component.board.message.MessageTableModel; import hae.instances.editor.RequestEditor; @@ -20,7 +20,7 @@ public class HaE implements BurpExtension { public void initialize(MontoyaApi api) { // 设置扩展名称 api.extension().setName("HaE - Highlighter and Extractor"); - String version = "4.0.5"; + String version = "4.0.6"; // 加载扩展后输出的项目信息 Logging logging = api.logging(); @@ -58,7 +58,7 @@ public class HaE implements BurpExtension { public void extensionUnloaded() { // 卸载清空数据 Config.globalDataMap.clear(); - CachePool.clear(); + MessageCache.clear(); } }); } diff --git a/src/main/java/hae/cache/DataQueryCache.java b/src/main/java/hae/cache/DataQueryCache.java new file mode 100644 index 0000000..85719eb --- /dev/null +++ b/src/main/java/hae/cache/DataQueryCache.java @@ -0,0 +1,46 @@ +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/CachePool.java b/src/main/java/hae/cache/MessageCache.java similarity index 97% rename from src/main/java/hae/cache/CachePool.java rename to src/main/java/hae/cache/MessageCache.java index eeacc5d..55973be 100644 --- a/src/main/java/hae/cache/CachePool.java +++ b/src/main/java/hae/cache/MessageCache.java @@ -6,7 +6,7 @@ import com.github.benmanes.caffeine.cache.Caffeine; import java.util.Map; import java.util.concurrent.TimeUnit; -public class CachePool { +public class MessageCache { private static final int MAX_SIZE = 100000; private static final int EXPIRE_DURATION = 5; diff --git a/src/main/java/hae/component/board/Databoard.java b/src/main/java/hae/component/board/Databoard.java index a9bfc06..043a21b 100644 --- a/src/main/java/hae/component/board/Databoard.java +++ b/src/main/java/hae/component/board/Databoard.java @@ -2,6 +2,7 @@ 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; @@ -23,19 +24,17 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.stream.Collectors; public class Databoard extends JPanel { + private static Boolean isMatchHost = false; private final MontoyaApi api; private final ConfigLoader configLoader; private final MessageTableModel messageTableModel; - + private final DefaultComboBoxModel comboBoxModel = new DefaultComboBoxModel(); + private final JComboBox hostComboBox = new JComboBox(comboBoxModel); private JTextField hostTextField; private JTabbedPane dataTabbedPane; private JSplitPane splitPane; private MessageTable messageTable; - - private static Boolean isMatchHost = false; - private final DefaultComboBoxModel comboBoxModel = new DefaultComboBoxModel(); - private final JComboBox hostComboBox = new JComboBox(comboBoxModel); - + private JProgressBar progressBar; private SwingWorker>, Void> handleComboBoxWorker; private SwingWorker applyHostFilterWorker; @@ -47,13 +46,25 @@ 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}; - ((GridBagLayout) getLayout()).rowHeights = new int[]{0, 65, 20, 0}; + ((GridBagLayout) getLayout()).rowHeights = new int[]{0, 65, 20, 0, 0}; ((GridBagLayout) getLayout()).columnWeights = new double[]{0.0, 0.0, 1.0, 0.0, 0.0, 1.0E-4}; - ((GridBagLayout) getLayout()).rowWeights = new double[]{0.0, 1.0, 0.0, 1.0E-4}; - + ((GridBagLayout) getLayout()).rowWeights = new double[]{0.0, 1.0, 0.0, 0.0, 1.0E-4}; JLabel hostLabel = new JLabel("Host:"); JButton clearButton = new JButton("Clear"); @@ -81,7 +92,7 @@ public class Databoard extends JPanel { clearButton.addActionListener(this::clearActionPerformed); - + progressBar = new JProgressBar(); splitPane.addComponentListener(new ComponentAdapter() { @Override public void componentResized(ComponentEvent e) { @@ -90,6 +101,7 @@ public class Databoard extends JPanel { }); splitPane.setVisible(false); + progressBar.setVisible(false); add(hostLabel, new GridBagConstraints(1, 0, 1, 1, 0.0, 0.0, GridBagConstraints.CENTER, GridBagConstraints.BOTH, new Insets(8, 0, 5, 5), 0, 0)); @@ -98,9 +110,12 @@ public class Databoard extends JPanel { add(actionButton, new GridBagConstraints(3, 0, 1, 1, 0.0, 0.0, GridBagConstraints.CENTER, GridBagConstraints.BOTH, new Insets(8, 0, 5, 5), 0, 0)); - add(splitPane, new GridBagConstraints(1, 1, 3, 2, 0.0, 1.0, + add(splitPane, new GridBagConstraints(1, 1, 3, 1, 0.0, 1.0, GridBagConstraints.CENTER, GridBagConstraints.BOTH, new Insets(0, 5, 0, 5), 0, 0)); + add(progressBar, new GridBagConstraints(1, 2, 3, 1, 1.0, 0.0, + GridBagConstraints.CENTER, GridBagConstraints.HORIZONTAL, + new Insets(0, 5, 0, 5), 0, 0)); hostComboBox.setMaximumRowCount(5); add(hostComboBox, new GridBagConstraints(2, 0, 1, 1, 0.0, 0.0, GridBagConstraints.CENTER, GridBagConstraints.BOTH, new Insets(8, 0, 5, 5), 0, 0)); @@ -120,6 +135,10 @@ public class Databoard extends JPanel { columnModel.getColumn(5).setPreferredWidth((int) (totalWidth * 0.1)); } + private void setProgressBar(boolean status) { + setProgressBar(status, progressBar, "Loading ..."); + } + private void setAutoMatch() { hostComboBox.setSelectedItem(null); hostComboBox.addActionListener(this::handleComboBoxAction); @@ -155,6 +174,8 @@ public class Databoard extends JPanel { String selectedHost = hostComboBox.getSelectedItem().toString(); if (getHostByList().contains(selectedHost)) { + progressBar.setVisible(true); + setProgressBar(true); hostTextField.setText(selectedHost); if (handleComboBoxWorker != null && !handleComboBoxWorker.isDone()) { @@ -193,6 +214,8 @@ public class Databoard extends JPanel { hostComboBox.setPopupVisible(false); applyHostFilter(selectedHost); + + setProgressBar(false); } } catch (Exception ignored) { } @@ -230,6 +253,12 @@ 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; @@ -255,6 +284,11 @@ public class Databoard extends JPanel { selectedDataMap = dataMap.get(selectedHost); } + // 将结果存入缓存 + if (selectedDataMap != null) { + DataQueryCache.putHostQueryResult(selectedHost, selectedDataMap); + } + return selectedDataMap; } @@ -314,19 +348,31 @@ public class Databoard extends JPanel { } private List getHostByList() { - if (!Config.globalDataMap.keySet().isEmpty()) { - return new ArrayList<>(Config.globalDataMap.keySet()); + // 先尝试从缓存获取结果 + List cachedResult = DataQueryCache.getHostFilterResult("all_hosts"); + if (cachedResult != null) { + return cachedResult; } - return new ArrayList<>(); + + 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(); if (retCode == JOptionPane.YES_OPTION && !host.isEmpty()) { dataTabbedPane.removeAll(); splitPane.setVisible(false); + progressBar.setVisible(false); Config.globalDataMap.keySet().parallelStream().forEach(key -> { if (StringProcessor.matchesHostPattern(key, host) || host.equals("*")) { @@ -353,8 +399,8 @@ public class Databoard extends JPanel { keysToRemove.forEach(Config.globalDataMap::remove); - if (Config.globalDataMap.keySet().size() == 1 && Config.globalDataMap.keySet().stream().anyMatch(key -> key.equals("*"))) { - Config.globalDataMap.keySet().remove("*"); + if (Config.globalDataMap.size() == 1 && Config.globalDataMap.keySet().stream().anyMatch(key -> key.equals("*"))) { + Config.globalDataMap.remove("*"); } messageTableModel.deleteByHost(host); diff --git a/src/main/java/hae/component/board/message/MessageTableModel.java b/src/main/java/hae/component/board/message/MessageTableModel.java index fdbf0de..d90db8c 100644 --- a/src/main/java/hae/component/board/message/MessageTableModel.java +++ b/src/main/java/hae/component/board/message/MessageTableModel.java @@ -10,7 +10,7 @@ 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.CachePool; +import hae.cache.MessageCache; import hae.utils.ConfigLoader; import hae.utils.DataManager; import hae.utils.string.HashCalculator; @@ -90,7 +90,7 @@ public class MessageTableModel extends AbstractTableModel { 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); @@ -98,7 +98,7 @@ public class MessageTableModel extends AbstractTableModel { splitPane.setRightComponent(messageTab); } - public void add(HttpRequestResponse messageInfo, String url, String method, String status, String length, String comment, String color, boolean flag) { + 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); @@ -157,6 +157,26 @@ public class MessageTableModel extends AbstractTableModel { } + 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]; + + // 复用现有的 add 方法逻辑，但跳过重复检查 + MessageEntry logEntry = new MessageEntry(messageInfo, method, url, comment, length, color, status); + log.add(logEntry); + } + } + // 批量更新完成后一次性通知表格更新 + fireTableDataChanged(); + } + public void deleteByHost(String filterText) { filteredLog.clear(); List rowsToRemove = new ArrayList<>(); @@ -317,7 +337,7 @@ public class MessageTableModel extends AbstractTableModel { private Map> getCacheData(byte[] content) { String hashIndex = HashCalculator.calculateHash(content); - return CachePool.get(hashIndex); + return MessageCache.get(hashIndex); } private boolean areMapsEqual(Map> map1, Map> map2) { @@ -426,11 +446,11 @@ public class MessageTableModel extends AbstractTableModel { } public class MessageTable extends JTable { - private MessageEntry messageEntry; private final ExecutorService executorService; - private int lastSelectedIndex = -1; private final HttpRequestEditor requestEditor; private final HttpResponseEditor responseEditor; + private MessageEntry messageEntry; + private int lastSelectedIndex = -1; public MessageTable(TableModel messageTableModel, HttpRequestEditor requestEditor, HttpResponseEditor responseEditor) { super(messageTableModel); diff --git a/src/main/java/hae/component/rule/Rules.java b/src/main/java/hae/component/rule/Rules.java index 0df6113..f2765ca 100644 --- a/src/main/java/hae/component/rule/Rules.java +++ b/src/main/java/hae/component/rule/Rules.java @@ -11,12 +11,41 @@ import java.awt.event.*; public class Rules extends JTabbedPane { private final MontoyaApi api; - private ConfigLoader configLoader; private final RuleProcessor ruleProcessor; private final JTextField ruleGroupNameTextField; - + private ConfigLoader configLoader; private Component tabComponent; private int selectedIndex; + private final Action cancelActionPerformed = new AbstractAction() { + @Override + public void actionPerformed(ActionEvent e) { + if (selectedIndex >= 0) { + setTabComponentAt(selectedIndex, tabComponent); + + ruleGroupNameTextField.setVisible(false); + ruleGroupNameTextField.setPreferredSize(null); + selectedIndex = -1; + tabComponent = null; + + requestFocusInWindow(); + } + } + }; + private final Action renameTitleActionPerformed = new AbstractAction() { + @Override + public void actionPerformed(ActionEvent e) { + String title = ruleGroupNameTextField.getText(); + if (!title.isEmpty() && selectedIndex >= 0) { + String oldName = getTitleAt(selectedIndex); + setTitleAt(selectedIndex, title); + + if (!oldName.equals(title)) { + ruleProcessor.renameRuleGroup(oldName, title); + } + } + cancelActionPerformed.actionPerformed(null); + } + }; public Rules(MontoyaApi api, ConfigLoader configLoader) { this.api = api; @@ -49,43 +78,48 @@ public class Rules extends JTabbedPane { addMouseListener(new MouseAdapter() { @Override public void mousePressed(MouseEvent e) { - int index = getSelectedIndex(); - Rectangle r = getBoundsAt(index); - if (r.contains(e.getPoint()) && index >= 0) { - switch (e.getButton()) { - case MouseEvent.BUTTON1: - if (e.getClickCount() == 2) { - selectedIndex = index; - tabComponent = getTabComponentAt(selectedIndex); - String ruleGroupName = getTitleAt(selectedIndex); + int index = indexAtLocation(e.getX(), e.getY()); + if (index < 0) { + return; + } - if (!"...".equals(ruleGroupName)) { - setTabComponentAt(selectedIndex, ruleGroupNameTextField); - ruleGroupNameTextField.setVisible(true); - ruleGroupNameTextField.setText(ruleGroupName); - ruleGroupNameTextField.selectAll(); - ruleGroupNameTextField.requestFocusInWindow(); - ruleGroupNameTextField.setMinimumSize(ruleGroupNameTextField.getPreferredSize()); - } - } else if (e.getClickCount() == 1) { - if ("...".equals(getTitleAt(getSelectedIndex()))) { - String title = ruleProcessor.newRule(); - Rule newRule = new Rule(api, configLoader, Config.ruleTemplate, tabbedPane); - insertTab(title, null, newRule, null, getTabCount() - 1); - setSelectedIndex(getTabCount() - 2); - } else { - renameTitleActionPerformed.actionPerformed(null); - } + switch (e.getButton()) { + case MouseEvent.BUTTON1: + if (e.getClickCount() == 2) { + selectedIndex = index; + tabComponent = getTabComponentAt(selectedIndex); + String ruleGroupName = getTitleAt(selectedIndex); + + if (!"...".equals(ruleGroupName)) { + setTabComponentAt(selectedIndex, ruleGroupNameTextField); + ruleGroupNameTextField.setVisible(true); + ruleGroupNameTextField.setText(ruleGroupName); + ruleGroupNameTextField.selectAll(); + ruleGroupNameTextField.requestFocusInWindow(); + ruleGroupNameTextField.setMinimumSize(ruleGroupNameTextField.getPreferredSize()); } - break; - case MouseEvent.BUTTON3: - if (!"...".equals(getTitleAt(getSelectedIndex()))) { - popupMenu.show(e.getComponent(), e.getX(), e.getY()); + } else if (e.getClickCount() == 1) { + String title = getTitleAt(index); + if ("...".equals(title)) { + // 阻止默认的选中行为 + e.consume(); + // 直接创建新标签 + String newTitle = ruleProcessor.newRule(); + Rule newRule = new Rule(api, configLoader, Config.ruleTemplate, Rules.this); + insertTab(newTitle, null, newRule, null, getTabCount() - 1); + setSelectedIndex(getTabCount() - 2); + } else { + renameTitleActionPerformed.actionPerformed(null); } - break; - default: - break; - } + } + break; + case MouseEvent.BUTTON3: + if (!"...".equals(getTitleAt(index))) { + popupMenu.show(e.getComponent(), e.getX(), e.getY()); + } + break; + default: + break; } } }); @@ -119,38 +153,6 @@ public class Rules extends JTabbedPane { } } } - - private final Action renameTitleActionPerformed = new AbstractAction() { - @Override - public void actionPerformed(ActionEvent e) { - String title = ruleGroupNameTextField.getText(); - if (!title.isEmpty() && selectedIndex >= 0) { - String oldName = getTitleAt(selectedIndex); - setTitleAt(selectedIndex, title); - - if (!oldName.equals(title)) { - ruleProcessor.renameRuleGroup(oldName, title); - } - } - cancelActionPerformed.actionPerformed(null); - } - }; - - private final Action cancelActionPerformed = new AbstractAction() { - @Override - public void actionPerformed(ActionEvent e) { - if (selectedIndex >= 0) { - setTabComponentAt(selectedIndex, tabComponent); - - ruleGroupNameTextField.setVisible(false); - ruleGroupNameTextField.setPreferredSize(null); - selectedIndex = -1; - tabComponent = null; - - requestFocusInWindow(); - } - } - }; } diff --git a/src/main/java/hae/instances/editor/RequestEditor.java b/src/main/java/hae/instances/editor/RequestEditor.java index 5749e9f..e025ca9 100644 --- a/src/main/java/hae/instances/editor/RequestEditor.java +++ b/src/main/java/hae/instances/editor/RequestEditor.java @@ -31,6 +31,28 @@ public class RequestEditor implements HttpRequestEditorProvider { this.configLoader = configLoader; } + public static boolean isListHasData(List> dataList) { + if (dataList != null && !dataList.isEmpty()) { + Map dataMap = dataList.get(0); + return dataMap != null && !dataMap.isEmpty(); + } + return false; + } + + public static void generateTabbedPaneFromResultMap(MontoyaApi api, ConfigLoader configLoader, JTabbedPane tabbedPane, List> result) { + tabbedPane.removeAll(); + if (result != null && !result.isEmpty()) { + Map dataMap = result.get(0); + if (dataMap != null && !dataMap.isEmpty()) { + dataMap.keySet().forEach(i -> { + String[] extractData = dataMap.get(i).split(Config.boundary); + Datatable dataPanel = new Datatable(api, configLoader, i, Arrays.asList(extractData)); + tabbedPane.addTab(i, dataPanel); + }); + } + } + } + @Override public ExtensionProvidedHttpRequestEditor provideHttpRequestEditor(EditorCreationContext editorCreationContext) { return new Editor(api, configLoader, editorCreationContext); @@ -42,11 +64,10 @@ public class RequestEditor implements HttpRequestEditorProvider { private final HttpUtils httpUtils; private final EditorCreationContext creationContext; private final MessageProcessor messageProcessor; + private final JTabbedPane jTabbedPane = new JTabbedPane(); private HttpRequestResponse requestResponse; private List> dataList; - private final JTabbedPane jTabbedPane = new JTabbedPane(); - public Editor(MontoyaApi api, ConfigLoader configLoader, EditorCreationContext creationContext) { this.api = api; this.configLoader = configLoader; @@ -118,26 +139,4 @@ public class RequestEditor implements HttpRequestEditorProvider { return false; } } - - public static boolean isListHasData(List> dataList) { - if (dataList != null && !dataList.isEmpty()) { - Map dataMap = dataList.get(0); - return dataMap != null && !dataMap.isEmpty(); - } - return false; - } - - public static void generateTabbedPaneFromResultMap(MontoyaApi api, ConfigLoader configLoader, JTabbedPane tabbedPane, List> result) { - tabbedPane.removeAll(); - if (result != null && !result.isEmpty()) { - Map dataMap = result.get(0); - if (dataMap != null && !dataMap.isEmpty()) { - dataMap.keySet().forEach(i -> { - String[] extractData = dataMap.get(i).split(Config.boundary); - Datatable dataPanel = new Datatable(api, configLoader, i, Arrays.asList(extractData)); - tabbedPane.addTab(i, dataPanel); - }); - } - } - } } diff --git a/src/main/java/hae/instances/editor/ResponseEditor.java b/src/main/java/hae/instances/editor/ResponseEditor.java index d8d4af2..5558c7a 100644 --- a/src/main/java/hae/instances/editor/ResponseEditor.java +++ b/src/main/java/hae/instances/editor/ResponseEditor.java @@ -41,11 +41,10 @@ public class ResponseEditor implements HttpResponseEditorProvider { private final HttpUtils httpUtils; private final EditorCreationContext creationContext; private final MessageProcessor messageProcessor; + private final JTabbedPane jTabbedPane = new JTabbedPane(); private HttpRequestResponse requestResponse; private List> dataList; - private final JTabbedPane jTabbedPane = new JTabbedPane(); - public Editor(MontoyaApi api, ConfigLoader configLoader, EditorCreationContext creationContext) { this.api = api; this.configLoader = configLoader; diff --git a/src/main/java/hae/instances/editor/WebSocketEditor.java b/src/main/java/hae/instances/editor/WebSocketEditor.java index a27dddd..13e93db 100644 --- a/src/main/java/hae/instances/editor/WebSocketEditor.java +++ b/src/main/java/hae/instances/editor/WebSocketEditor.java @@ -36,11 +36,10 @@ public class WebSocketEditor implements WebSocketMessageEditorProvider { private final ConfigLoader configLoader; private final EditorCreationContext creationContext; private final MessageProcessor messageProcessor; + private final JTabbedPane jTabbedPane = new JTabbedPane(); private ByteArray message; private List> dataList; - private final JTabbedPane jTabbedPane = new JTabbedPane(); - public Editor(MontoyaApi api, ConfigLoader configLoader, EditorCreationContext creationContext) { this.api = api; this.configLoader = configLoader; diff --git a/src/main/java/hae/instances/http/utils/MessageProcessor.java b/src/main/java/hae/instances/http/utils/MessageProcessor.java index 84b7524..16ba4c7 100644 --- a/src/main/java/hae/instances/http/utils/MessageProcessor.java +++ b/src/main/java/hae/instances/http/utils/MessageProcessor.java @@ -10,7 +10,6 @@ import java.nio.charset.StandardCharsets; import java.util.*; import java.util.stream.Collectors; - public class MessageProcessor { private final MontoyaApi api; private final RegularMatcher regularMatcher; diff --git a/src/main/java/hae/instances/http/utils/RegularMatcher.java b/src/main/java/hae/instances/http/utils/RegularMatcher.java index cce2764..953313e 100644 --- a/src/main/java/hae/instances/http/utils/RegularMatcher.java +++ b/src/main/java/hae/instances/http/utils/RegularMatcher.java @@ -8,7 +8,7 @@ import dk.brics.automaton.AutomatonMatcher; import dk.brics.automaton.RegExp; import dk.brics.automaton.RunAutomaton; import hae.Config; -import hae.cache.CachePool; +import hae.cache.MessageCache; import hae.utils.DataManager; import hae.utils.string.HashCalculator; import hae.utils.string.StringProcessor; @@ -27,10 +27,57 @@ public class RegularMatcher { } + public synchronized static void putDataToGlobalMap(MontoyaApi api, String host, String name, List dataList, boolean flag) { + // 添加到全局变量中，便于Databoard检索 + if (!Objects.equals(host, "") && host != null) { + Config.globalDataMap.compute(host, (existingHost, existingMap) -> { + Map> gRuleMap = Optional.ofNullable(existingMap).orElse(new ConcurrentHashMap<>()); + + gRuleMap.merge(name, new ArrayList<>(dataList), (existingList, newList) -> { + Set combinedSet = new LinkedHashSet<>(existingList); + combinedSet.addAll(newList); + return new ArrayList<>(combinedSet); + }); + + if (flag) { + // 数据存储在BurpSuite空间内 + try { + DataManager dataManager = new DataManager(api); + PersistedObject persistedObject = PersistedObject.persistedObject(); + gRuleMap.forEach((kName, vList) -> { + PersistedList persistedList = PersistedList.persistedStringList(); + persistedList.addAll(vList); + persistedObject.setStringList(kName, persistedList); + }); + dataManager.putData("data", host, persistedObject); + } catch (Exception ignored) { + } + } + + return gRuleMap; + }); + + String[] splitHost = host.split("\\."); + String onlyHost = host.split(":")[0]; + + String anyHost = (splitHost.length > 2 && !StringProcessor.matchHostIsIp(onlyHost)) ? StringProcessor.replaceFirstOccurrence(onlyHost, splitHost[0], "*") : ""; + + if (!Config.globalDataMap.containsKey(anyHost) && !anyHost.isEmpty()) { + // 添加通配符Host，实际数据从查询哪里将所有数据提取 + Config.globalDataMap.put(anyHost, new HashMap<>()); + } + + if (!Config.globalDataMap.containsKey("*")) { + // 添加通配符全匹配，同上 + Config.globalDataMap.put("*", new HashMap<>()); + } + } + } + public Map> match(String host, String type, String message, String header, String body) { // 先从缓存池里判断是否有已经匹配好的结果 String messageIndex = HashCalculator.calculateHash(message.getBytes()); - Map> map = CachePool.get(messageIndex); + Map> map = MessageCache.get(messageIndex); if (map != null) { return map; } else { @@ -106,58 +153,11 @@ public class RegularMatcher { } } }); - CachePool.put(messageIndex, finalMap); + MessageCache.put(messageIndex, finalMap); return finalMap; } } - public synchronized static void putDataToGlobalMap(MontoyaApi api, String host, String name, List dataList, boolean flag) { - // 添加到全局变量中，便于Databoard检索 - if (!Objects.equals(host, "") && host != null) { - Config.globalDataMap.compute(host, (existingHost, existingMap) -> { - Map> gRuleMap = Optional.ofNullable(existingMap).orElse(new ConcurrentHashMap<>()); - - gRuleMap.merge(name, new ArrayList<>(dataList), (existingList, newList) -> { - Set combinedSet = new LinkedHashSet<>(existingList); - combinedSet.addAll(newList); - return new ArrayList<>(combinedSet); - }); - - if (flag) { - // 数据存储在BurpSuite空间内 - try { - DataManager dataManager = new DataManager(api); - PersistedObject persistedObject = PersistedObject.persistedObject(); - gRuleMap.forEach((kName, vList) -> { - PersistedList persistedList = PersistedList.persistedStringList(); - persistedList.addAll(vList); - persistedObject.setStringList(kName, persistedList); - }); - dataManager.putData("data", host, persistedObject); - } catch (Exception ignored) { - } - } - - return gRuleMap; - }); - - String[] splitHost = host.split("\\."); - String onlyHost = host.split(":")[0]; - - String anyHost = (splitHost.length > 2 && !StringProcessor.matchHostIsIp(onlyHost)) ? StringProcessor.replaceFirstOccurrence(onlyHost, splitHost[0], "*") : ""; - - if (!Config.globalDataMap.containsKey(anyHost) && !anyHost.isEmpty()) { - // 添加通配符Host，实际数据从查询哪里将所有数据提取 - Config.globalDataMap.put(anyHost, new HashMap<>()); - } - - if (!Config.globalDataMap.containsKey("*")) { - // 添加通配符全匹配，同上 - Config.globalDataMap.put("*", new HashMap<>()); - } - } - } - private List matchByRegex(String f_regex, String s_regex, String content, String format, String engine, boolean sensitive) { List retList = new ArrayList<>(); if ("nfa".equals(engine)) { diff --git a/src/main/java/hae/utils/ConfigLoader.java b/src/main/java/hae/utils/ConfigLoader.java index 4cb61a6..86228a1 100644 --- a/src/main/java/hae/utils/ConfigLoader.java +++ b/src/main/java/hae/utils/ConfigLoader.java @@ -49,6 +49,11 @@ public class ConfigLoader { Config.globalRules = getRules(); } + private static boolean isValidConfigPath(String configPath) { + File configPathFile = new File(configPath); + return configPathFile.exists() && configPathFile.isDirectory(); + } + private String determineConfigPath() { // 优先级1：用户根目录 String userConfigPath = String.format("%s/.config/HaE", System.getProperty("user.home")); @@ -67,11 +72,6 @@ public class ConfigLoader { return userConfigPath; } - private static boolean isValidConfigPath(String configPath) { - File configPathFile = new File(configPath); - return configPathFile.exists() && configPathFile.isDirectory(); - } - public void initConfig() { Map r = new LinkedHashMap<>(); r.put("ExcludeSuffix", getExcludeSuffix()); @@ -102,8 +102,6 @@ public class ConfigLoader { Representer representer = new Representer(dop); Map rulesMap = new Yaml(representer, dop).load(inputStream); - String[] fieldKeys = {"loaded", "name", "f_regex", "s_regex", "format", "color", "scope", "engine", "sensitive"}; - Object rulesObj = rulesMap.get("rules"); if (rulesObj instanceof List) { List> groupData = (List>) rulesObj; @@ -114,9 +112,9 @@ public class ConfigLoader { if (ruleObj instanceof List) { List> ruleData = (List>) ruleObj; for (Map ruleFields : ruleData) { - Object[] valuesArray = new Object[fieldKeys.length]; - for (int i = 0; i < fieldKeys.length; i++) { - valuesArray[i] = ruleFields.get(fieldKeys[i]); + Object[] valuesArray = new Object[Config.ruleFields.length]; + for (int i = 0; i < Config.ruleFields.length; i++) { + valuesArray[i] = ruleFields.get(Config.ruleFields[i].toLowerCase().replace("-", "_")); } data.add(valuesArray); } @@ -138,26 +136,50 @@ public class ConfigLoader { return getValueFromConfig("BlockHost", Config.host); } + public void setBlockHost(String blockHost) { + setValueToConfig("BlockHost", blockHost); + } + public String getExcludeSuffix() { return getValueFromConfig("ExcludeSuffix", Config.suffix); } + public void setExcludeSuffix(String excludeSuffix) { + setValueToConfig("ExcludeSuffix", excludeSuffix); + } + public String getExcludeStatus() { return getValueFromConfig("ExcludeStatus", Config.status); } + public void setExcludeStatus(String status) { + setValueToConfig("ExcludeStatus", status); + } + public String getLimitSize() { return getValueFromConfig("LimitSize", Config.size); } + public void setLimitSize(String size) { + setValueToConfig("LimitSize", size); + } + public String getScope() { return getValueFromConfig("HaEScope", Config.scopeOptions); } + public void setScope(String scope) { + setValueToConfig("HaEScope", scope); + } + public boolean getMode() { return getValueFromConfig("HaEModeStatus", Config.modeStatus).equals("true"); } + public void setMode(String mode) { + setValueToConfig("HaEModeStatus", mode); + } + private String getValueFromConfig(String name, String defaultValue) { File yamlSetting = new File(configFilePath); if (!yamlSetting.exists() || !yamlSetting.isFile()) { @@ -176,30 +198,6 @@ public class ConfigLoader { return defaultValue; } - public void setExcludeSuffix(String excludeSuffix) { - setValueToConfig("ExcludeSuffix", excludeSuffix); - } - - public void setBlockHost(String blockHost) { - setValueToConfig("BlockHost", blockHost); - } - - public void setExcludeStatus(String status) { - setValueToConfig("ExcludeStatus", status); - } - - public void setLimitSize(String size) { - setValueToConfig("LimitSize", size); - } - - public void setScope(String scope) { - setValueToConfig("HaEScope", scope); - } - - public void setMode(String mode) { - setValueToConfig("HaEModeStatus", mode); - } - private void setValueToConfig(String name, String value) { Map currentConfig = loadCurrentConfig(); currentConfig.put(name, value); diff --git a/src/main/java/hae/utils/DataManager.java b/src/main/java/hae/utils/DataManager.java index 894acfa..0749ab2 100644 --- a/src/main/java/hae/utils/DataManager.java +++ b/src/main/java/hae/utils/DataManager.java @@ -10,6 +10,12 @@ 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.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; + public class DataManager { private final MontoyaApi api; private final Persistence persistence; @@ -69,22 +75,81 @@ public class DataManager { } private void loadMessageData(PersistedList messageIndex, MessageTableModel messageTableModel) { - if (messageIndex != null && !messageIndex.isEmpty()) { - messageIndex.forEach(index -> { + if (messageIndex == null || messageIndex.isEmpty()) { + 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()); + } + } + + final int batchSize = 2000; // 增加批处理大小 + final int threadCount = Math.max(8, Runtime.getRuntime().availableProcessors() * 2); // 增加线程数 + int totalSize = indexList.size(); + + // 使用更高效的线程池 + 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); + } + } catch (Exception e) { + api.logging().logToError("批量添加数据时出错: " + e.getMessage()); + } finally { + executorService.shutdown(); + } + } + + private List processBatchParallel(List batch) { + List batchData = new ArrayList<>(); + for (String index : batch) { + try { PersistedObject dataObj = persistence.extensionData().getChildObject(index); if (dataObj != null) { HttpRequestResponse messageInfo = dataObj.getHttpRequestResponse("messageInfo"); - String comment = dataObj.getString("comment"); - String color = dataObj.getString("color"); - HttpRequest request = messageInfo.request(); - HttpResponse response = messageInfo.response(); - String method = request.method(); - String url = request.url(); - String status = String.valueOf(response.statusCode()); - String length = String.valueOf(response.toByteArray().length()); - messageTableModel.add(messageInfo, url, method, status, length, comment, color, false); + if (messageInfo != null) { + batchData.add(prepareMessageData(messageInfo, dataObj)); + } } - }); + } catch (Exception e) { + api.logging().logToError("处理消息数据时出错: " + e.getMessage() + ", index: " + index); + } } + return batchData; } -} + + private Object[] prepareMessageData(HttpRequestResponse messageInfo, PersistedObject dataObj) { + HttpRequest request = messageInfo.request(); + HttpResponse response = messageInfo.response(); + return new Object[]{ + messageInfo, + request.url(), + request.method(), + String.valueOf(response.statusCode()), + String.valueOf(response.toByteArray().length()), + dataObj.getString("comment"), + dataObj.getString("color"), + false + }; + } +} \ No newline at end of file