Skip to content

Commit

Permalink
Version: 3.2.2 Update
Browse files Browse the repository at this point in the history
  • Loading branch information
@gh0stkey
gh0stkey committed Jun 19, 2024
1 parent 54973d9 commit e5f55b6
Show file tree
Hide file tree
Showing 9 changed files with 192 additions and 192 deletions.
13 changes: 12 additions & 1 deletion src/main/java/hae/HaE.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,9 @@

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.component.Main;
import hae.component.board.message.MessageTableModel;
import hae.instances.editor.RequestEditor;
Expand All @@ -16,7 +18,7 @@ public class HaE implements BurpExtension {
@Override
public void initialize(MontoyaApi api) {
// 设置扩展名称
String version = "3.2.1";
String version = "3.2.2";
api.extension().setName(String.format("HaE (%s) - Highlighter and Extractor", version));

// 加载扩展后输出的项目信息
Expand All @@ -43,5 +45,14 @@ public void initialize(MontoyaApi api) {
api.userInterface().registerHttpRequestEditorProvider(new RequestEditor(api, configLoader));
api.userInterface().registerHttpResponseEditorProvider(new ResponseEditor(api, configLoader));
api.userInterface().registerWebSocketMessageEditorProvider(new WebSocketEditor(api));

api.extension().registerUnloadingHandler(new ExtensionUnloadingHandler() {
@Override
public void extensionUnloaded() {
// 卸载清空数据
Config.globalDataMap.clear();
CachePool.clear();
}
});
}
}
151 changes: 89 additions & 62 deletions src/main/java/hae/component/board/Databoard.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,10 @@
import java.io.File;
import java.util.List;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.Function;
import java.util.stream.Collectors;

Expand All @@ -43,7 +46,10 @@ public class Databoard extends JPanel {
private final DefaultComboBoxModel comboBoxModel = new DefaultComboBoxModel();
private final JComboBox hostComboBox = new JComboBox(comboBoxModel);

private SwingWorker<Boolean, Void> currentWorker;
private SwingWorker<Boolean, Void> handleComboBoxWorker;
private SwingWorker<Void, Void> applyHostFilterWorker;
private SwingWorker<List<String>, Void> exportActionWorker;
private SwingWorker<List<String>, Void> importActionWorker;

public Databoard(MontoyaApi api, ConfigLoader configLoader, MessageTableModel messageTableModel) {
this.api = api;
Expand Down Expand Up @@ -77,11 +83,10 @@ private void initComponents() {

hostTextField = new JTextField();
splitPane = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT);
dataTabbedPane = new JTabbedPane(JTabbedPane.TOP);

dataTabbedPane = new JTabbedPane(JTabbedPane.TOP);
dataTabbedPane.setPreferredSize(new Dimension(500, 0));
dataTabbedPane.setTabLayoutPolicy(JTabbedPane.SCROLL_TAB_LAYOUT);
splitPane.setLeftComponent(dataTabbedPane);

actionButton.addActionListener(e -> {
int x = 0;
Expand Down Expand Up @@ -185,42 +190,46 @@ private void handleComboBoxAction(ActionEvent e) {
progressBar.setVisible(true);
setProgressBar(true);
String selectedHost = hostComboBox.getSelectedItem().toString();
hostTextField.setText(selectedHost);

if (currentWorker != null && !currentWorker.isDone()) {
currentWorker.cancel(true);
}
if (getHostByList().contains(selectedHost)) {
hostTextField.setText(selectedHost);

currentWorker = new SwingWorker<Boolean, Void>() {
@Override
protected Boolean doInBackground() {
return populateTabbedPaneByHost(selectedHost);
if (handleComboBoxWorker != null && !handleComboBoxWorker.isDone()) {
handleComboBoxWorker.cancel(true);
}

@Override
protected void done() {
if (!isCancelled()) {
try {
boolean status = get();
if (status) {
JSplitPane messageSplitPane = messageTableModel.getSplitPane();
splitPane.setRightComponent(messageSplitPane);
messageTable = messageTableModel.getMessageTable();
resizePanel();

splitPane.setVisible(true);
hostTextField.setText(selectedHost);

hostComboBox.setPopupVisible(false);
applyHostFilter(selectedHost);
handleComboBoxWorker = new SwingWorker<Boolean, Void>() {
@Override
protected Boolean doInBackground() {
return populateTabbedPaneByHost(selectedHost);
}

@Override
protected void done() {
if (!isCancelled()) {
try {
boolean status = get();
if (status) {
JSplitPane messageSplitPane = messageTableModel.getSplitPane();
splitPane.setLeftComponent(dataTabbedPane);
splitPane.setRightComponent(messageSplitPane);
messageTable = messageTableModel.getMessageTable();
resizePanel();

splitPane.setVisible(true);
hostTextField.setText(selectedHost);

hostComboBox.setPopupVisible(false);
applyHostFilter(selectedHost);
}
} catch (Exception ignored) {
}
} catch (Exception ignored) {
}
}
}
};
};

currentWorker.execute();
handleComboBoxWorker.execute();
}
}
}

Expand Down Expand Up @@ -254,7 +263,7 @@ private boolean populateTabbedPaneByHost(String selectedHost) {

if (selectedHost.contains("*")) {
selectedDataMap = new HashMap<>();
dataMap.keySet().parallelStream().forEach(key -> {
dataMap.keySet().forEach(key -> {
if ((StringProcessor.matchesHostPattern(key, selectedHost) || selectedHost.equals("*")) && !key.contains("*")) {
Map<String, List<String>> ruleMap = dataMap.get(key);
for (String ruleKey : ruleMap.keySet()) {
Expand Down Expand Up @@ -286,6 +295,7 @@ private boolean populateTabbedPaneByHost(String selectedHost) {

return true;
}

return false;
}

Expand Down Expand Up @@ -316,7 +326,11 @@ private void applyHostFilter(String filterText) {
TableRowSorter<TableModel> sorter = (TableRowSorter<TableModel>) messageTable.getRowSorter();
String cleanedText = StringProcessor.replaceFirstOccurrence(filterText, "*.", "");

new SwingWorker<Void, Void>() {
if (applyHostFilterWorker != null && !applyHostFilterWorker.isDone()) {
applyHostFilterWorker.cancel(true);
}

applyHostFilterWorker = new SwingWorker<Void, Void>() {
@Override
protected Void doInBackground() throws Exception {
RowFilter<Object, Object> rowFilter = new RowFilter<Object, Object>() {
Expand All @@ -340,8 +354,9 @@ public boolean include(Entry<?, ?> entry) {
protected void done() {
setProgressBar(false);
}
}.execute();
};

applyHostFilterWorker.execute();
}

private List<String> getHostByList() {
Expand All @@ -364,7 +379,11 @@ private void exportActionPerformed(ActionEvent e) {
return;
}

new SwingWorker<List<String>, Void>() {
if (exportActionWorker != null && !exportActionWorker.isDone()) {
exportActionWorker.cancel(true);
}

exportActionWorker = new SwingWorker<List<String>, Void>() {
@Override
protected List<String> doInBackground() {
ConcurrentHashMap<String, Map<String, List<String>>> dataMap = Config.globalDataMap;
Expand All @@ -382,7 +401,9 @@ protected void done() {
} catch (Exception ignored) {
}
}
}.execute();
};

exportActionWorker.execute();
}

private List<String> exportData(String selectedHost, String exportDir, Map<String, Map<String, List<String>>> dataMap) {
Expand Down Expand Up @@ -469,7 +490,11 @@ private void importActionPerformed(ActionEvent e) {
return;
}

new SwingWorker<List<String>, Void>() {
if (importActionWorker != null && !importActionWorker.isDone()) {
importActionWorker.cancel(true);
}

importActionWorker = new SwingWorker<List<String>, Void>() {
@Override
protected List<String> doInBackground() {
List<String> filesWithExtension = findFilesWithExtension(new File(exportDir), ".hae");
Expand All @@ -489,23 +514,26 @@ protected void done() {
} catch (Exception ignored) {
}
}
}.execute();
};

importActionWorker.execute();
}

private String importData(String filename) {
ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2);

HaeFileContent haeFileContent = projectProcessor.readHaeFile(filename);
boolean readStatus = haeFileContent != null;

ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2);
List<Future<?>> futures = new ArrayList<>();
List<Callable<Void>> tasks = new ArrayList<>();

if (readStatus) {
try {
String host = haeFileContent.getHost();
haeFileContent.getDataMap().forEach((key, value) -> RegularMatcher.putDataToGlobalMap(host, key, value));

haeFileContent.getUrlMap().forEach((key, urlItemMap) -> {
Future<?> future = executor.submit(() -> {
tasks.add(() -> {
String url = urlItemMap.get("url");
String comment = urlItemMap.get("comment");
String color = urlItemMap.get("color");
Expand All @@ -515,18 +543,11 @@ private String importData(String filename) {
String path = haeFileContent.getHttpPath();

messageTableModel.add(null, url, method, status, length, comment, color, key, path);
return null;
});

futures.add(future);
});

for (Future<?> future : futures) {
try {
future.get();
} catch (InterruptedException | ExecutionException e) {
Thread.currentThread().interrupt();
}
}
executor.invokeAll(tasks);
} catch (Exception e) {
api.logging().logToError("importData: " + e.getMessage());
} finally {
Expand Down Expand Up @@ -590,20 +611,26 @@ private void clearActionPerformed(ActionEvent e) {
}
});

if (!StringProcessor.matchHostIsIp(host) && !host.contains("*.")) {
String baseDomain = StringProcessor.getBaseDomain(StringProcessor.extractHostname(host));
// 删除无用的数据
Set<String> wildcardKeys = Config.globalDataMap.keySet().stream()
.filter(key -> key.startsWith("*."))
.collect(Collectors.toSet());

long count = Config.globalDataMap.keySet().stream()
.filter(k -> !k.equals("*." + baseDomain))
.filter(k -> StringProcessor.matchFromEnd(k, baseDomain))
.count();
Set<String> existingSuffixes = Config.globalDataMap.keySet().stream()
.filter(key -> !key.startsWith("*."))
.map(key -> {
int dotIndex = key.indexOf(".");
return dotIndex != -1 ? key.substring(dotIndex) : "";
})
.collect(Collectors.toSet());

if (count == 0) {
Config.globalDataMap.remove("*." + baseDomain);
}
}
Set<String> keysToRemove = wildcardKeys.stream()
.filter(key -> !existingSuffixes.contains(key.substring(1)))
.collect(Collectors.toSet());

keysToRemove.forEach(Config.globalDataMap::remove);

if (Config.globalDataMap.keySet().size() == 1 && Config.globalDataMap.keySet().stream().anyMatch(key -> key.contains("*"))) {
if (Config.globalDataMap.keySet().size() == 1 && Config.globalDataMap.keySet().stream().anyMatch(key -> key.equals("*"))) {
Config.globalDataMap.keySet().remove("*");
}

Expand Down
8 changes: 3 additions & 5 deletions src/main/java/hae/component/board/Datatable.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ public void setTableListener(MessageTableModel messagePanel) {
@Override
public void exportToClipboard(JComponent comp, Clipboard clip, int action) throws IllegalStateException {
if (comp instanceof JTable) {
StringSelection stringSelection = new StringSelection(getSelectedDataAtTable((JTable) comp));
StringSelection stringSelection = new StringSelection(getSelectedDataAtTable((JTable) comp).replace("\0", "").replaceAll("[\\p{Cntrl}&&[^\r\n\t]]", ""));
clip.setContents(stringSelection, null);
} else {
super.exportToClipboard(comp, clip, action);
Expand Down Expand Up @@ -191,13 +191,11 @@ public String getSelectedDataAtTable(JTable table) {
selectData.append(table.getValueAt(row, 1).toString()).append("\n");
}

// 便于单行复制，去除最后一个换行符
if (!selectData.isEmpty()) {
selectData.deleteCharAt(selectData.length() - 1);
return selectData.toString();
} else {
return "";
}

return selectData.toString();
}

public JTable getDataTable() {
Expand Down
6 changes: 3 additions & 3 deletions src/main/java/hae/component/board/message/MessageRenderer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,16 @@
import javax.swing.table.DefaultTableCellRenderer;
import java.awt.*;
import java.util.HashMap;
import java.util.List;
import java.util.LinkedList;
import java.util.Map;

public class MessageRenderer extends DefaultTableCellRenderer {

private final List<MessageEntry> log;
private final LinkedList<MessageEntry> log;
private final Map<String, Color> colorMap = new HashMap<>();
private final JTable table; // 保存对表格的引用

public MessageRenderer(List<MessageEntry> log, JTable table) {
public MessageRenderer(LinkedList<MessageEntry> log, JTable table) {
this.log = log;
// 与BurpSuite的颜色保持一致
this.colorMap.put("red", new Color(0xFF, 0x64, 0x64));
Expand Down
Loading

0 comments on commit e5f55b6

Please sign in to comment.