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

Version: 4.1 Update

Browse Source
This commit is contained in:
gh0stkey
2025-03-21 21:22:11 +08:00
parent 1a5ed2a6a3
commit 20afa30822
13 changed files with 393 additions and 220 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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();
}
});
}

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

@@ -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<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();
}
}

2
src/main/java/hae/cache/CachePool.java → src/main/java/hae/cache/MessageCache.java vendored
View File

@@ -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;

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

@@ -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<Map<String, List<String>>, Void> handleComboBoxWorker;
private SwingWorker<Void, Void> 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<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;
@@ -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<String> getHostByList() {
if (!Config.globalDataMap.keySet().isEmpty()) {
return new ArrayList<>(Config.globalDataMap.keySet());
// 先尝试从缓存获取结果
List<String> cachedResult = DataQueryCache.getHostFilterResult("all_hosts");
if (cachedResult != null) {
return cachedResult;
}
return new ArrayList<>();
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();
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);

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

@@ -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<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];
// 复用现有的 add 方法逻辑，但跳过重复检查
MessageEntry logEntry = new MessageEntry(messageInfo, method, url, comment, length, color, status);
log.add(logEntry);
}
}
// 批量更新完成后一次性通知表格更新
fireTableDataChanged();
}
public void deleteByHost(String filterText) {
filteredLog.clear();
List<Integer> rowsToRemove = new ArrayList<>();
@@ -317,7 +337,7 @@ public class MessageTableModel extends AbstractTableModel {
private Map<String, Map<String, Object>> getCacheData(byte[] content) {
String hashIndex = HashCalculator.calculateHash(content);
return CachePool.get(hashIndex);
return MessageCache.get(hashIndex);
}
private boolean areMapsEqual(Map<String, Map<String, Object>> map1, Map<String, Map<String, Object>> 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);

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

@@ -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();
}
}
};
}

47
src/main/java/hae/instances/editor/RequestEditor.java
View File

@@ -31,6 +31,28 @@ public class RequestEditor implements HttpRequestEditorProvider {
this.configLoader = configLoader;
}
public static boolean isListHasData(List<Map<String, String>> dataList) {
if (dataList != null && !dataList.isEmpty()) {
Map<String, String> dataMap = dataList.get(0);
return dataMap != null && !dataMap.isEmpty();
}
return false;
}
public static void generateTabbedPaneFromResultMap(MontoyaApi api, ConfigLoader configLoader, JTabbedPane tabbedPane, List<Map<String, String>> result) {
tabbedPane.removeAll();
if (result != null && !result.isEmpty()) {
Map<String, String> 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<Map<String, String>> 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<Map<String, String>> dataList) {
if (dataList != null && !dataList.isEmpty()) {
Map<String, String> dataMap = dataList.get(0);
return dataMap != null && !dataMap.isEmpty();
}
return false;
}
public static void generateTabbedPaneFromResultMap(MontoyaApi api, ConfigLoader configLoader, JTabbedPane tabbedPane, List<Map<String, String>> result) {
tabbedPane.removeAll();
if (result != null && !result.isEmpty()) {
Map<String, String> 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);
});
}
}
}
}

3
src/main/java/hae/instances/editor/ResponseEditor.java
View File

@@ -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<Map<String, String>> dataList;
private final JTabbedPane jTabbedPane = new JTabbedPane();
public Editor(MontoyaApi api, ConfigLoader configLoader, EditorCreationContext creationContext) {
this.api = api;
this.configLoader = configLoader;

3
src/main/java/hae/instances/editor/WebSocketEditor.java
View File

@@ -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<Map<String, String>> dataList;
private final JTabbedPane jTabbedPane = new JTabbedPane();
public Editor(MontoyaApi api, ConfigLoader configLoader, EditorCreationContext creationContext) {
this.api = api;
this.configLoader = configLoader;

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

@@ -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;

100
src/main/java/hae/instances/http/utils/RegularMatcher.java
View File

@@ -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<String> dataList, boolean flag) {
// 添加到全局变量中，便于Databoard检索
if (!Objects.equals(host, "") && host != null) {
Config.globalDataMap.compute(host, (existingHost, existingMap) -> {
Map<String, List<String>> gRuleMap = Optional.ofNullable(existingMap).orElse(new ConcurrentHashMap<>());
gRuleMap.merge(name, new ArrayList<>(dataList), (existingList, newList) -> {
Set<String> 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<String> 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<String, Map<String, Object>> match(String host, String type, String message, String header, String body) {
// 先从缓存池里判断是否有已经匹配好的结果
String messageIndex = HashCalculator.calculateHash(message.getBytes());
Map<String, Map<String, Object>> map = CachePool.get(messageIndex);
Map<String, Map<String, Object>> 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<String> dataList, boolean flag) {
// 添加到全局变量中，便于Databoard检索
if (!Objects.equals(host, "") && host != null) {
Config.globalDataMap.compute(host, (existingHost, existingMap) -> {
Map<String, List<String>> gRuleMap = Optional.ofNullable(existingMap).orElse(new ConcurrentHashMap<>());
gRuleMap.merge(name, new ArrayList<>(dataList), (existingList, newList) -> {
Set<String> 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<String> 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<String> matchByRegex(String f_regex, String s_regex, String content, String format, String engine, boolean sensitive) {
List<String> retList = new ArrayList<>();
if ("nfa".equals(engine)) {

66
src/main/java/hae/utils/ConfigLoader.java
View File

@@ -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<String, Object> r = new LinkedHashMap<>();
r.put("ExcludeSuffix", getExcludeSuffix());
@@ -102,8 +102,6 @@ public class ConfigLoader {
Representer representer = new Representer(dop);
Map<String, Object> 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<Map<String, Object>> groupData = (List<Map<String, Object>>) rulesObj;
@@ -114,9 +112,9 @@ public class ConfigLoader {
if (ruleObj instanceof List) {
List<Map<String, Object>> ruleData = (List<Map<String, Object>>) ruleObj;
for (Map<String, Object> 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<String, Object> currentConfig = loadCurrentConfig();
currentConfig.put(name, value);

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

@@ -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<String> messageIndex, MessageTableModel messageTableModel) {
if (messageIndex != null && !messageIndex.isEmpty()) {
messageIndex.forEach(index -> {
if (messageIndex == null || messageIndex.isEmpty()) {
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());
}
}
final int batchSize = 2000; // 增加批处理大小
final int threadCount = Math.max(8, Runtime.getRuntime().availableProcessors() * 2); // 增加线程数
int totalSize = indexList.size();
// 使用更高效的线程池
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);
}
} 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) {
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
};
}
}