Skip to content

Commit

Permalink
Version: 3.3.1 Update
Browse files Browse the repository at this point in the history
  • Loading branch information
@gh0stkey
gh0stkey committed Aug 12, 2024
1 parent a7e0a2a commit 4f04013
Show file tree
Hide file tree
Showing 11 changed files with 198 additions and 134 deletions.
2 changes: 2 additions & 0 deletions src/main/java/hae/Config.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@ public class Config {

public static String host = "gh0st.cn";

public static String status = "404";

public static String[] scope = new String[]{
"any",
"any header",
Expand Down
2 changes: 1 addition & 1 deletion src/main/java/hae/HaE.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ public class HaE implements BurpExtension {
@Override
public void initialize(MontoyaApi api) {
// 设置扩展名称
String version = "3.3";
String version = "3.3.1";
api.extension().setName(String.format("HaE (%s) - Highlighter and Extractor", version));

// 加载扩展后输出的项目信息
Expand Down
74 changes: 47 additions & 27 deletions src/main/java/hae/component/board/Databoard.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,11 @@
import hae.utils.string.StringProcessor;

import javax.swing.*;
import javax.swing.border.TitledBorder;
import javax.swing.event.DocumentEvent;
import javax.swing.event.DocumentListener;
import javax.swing.filechooser.FileNameExtensionFilter;
import javax.swing.table.DefaultTableModel;
import javax.swing.table.TableColumnModel;
import javax.swing.table.TableModel;
import javax.swing.table.TableRowSorter;
Expand Down Expand Up @@ -49,8 +51,8 @@ public class Databoard extends JPanel {

private SwingWorker<Map<String, List<String>>, Void> handleComboBoxWorker;
private SwingWorker<Void, Void> applyHostFilterWorker;
private SwingWorker<List<String>, Void> exportActionWorker;
private SwingWorker<List<String>, Void> importActionWorker;
private SwingWorker<List<Object[]>, Void> exportActionWorker;
private SwingWorker<List<Object[]>, Void> importActionWorker;

public Databoard(MontoyaApi api, ConfigLoader configLoader, MessageTableModel messageTableModel) {
this.api = api;
Expand Down Expand Up @@ -193,11 +195,11 @@ public void changedUpdate(DocumentEvent e) {

private void handleComboBoxAction(ActionEvent e) {
if (!isMatchHost && hostComboBox.getSelectedItem() != null) {
progressBar.setVisible(true);
setProgressBar(true);
String selectedHost = hostComboBox.getSelectedItem().toString();

if (getHostByList().contains(selectedHost)) {
progressBar.setVisible(true);
setProgressBar(true);
hostTextField.setText(selectedHost);

if (handleComboBoxWorker != null && !handleComboBoxWorker.isDone()) {
Expand Down Expand Up @@ -385,21 +387,19 @@ private void exportActionPerformed(ActionEvent e) {
exportActionWorker.cancel(true);
}

exportActionWorker = new SwingWorker<List<String>, Void>() {
exportActionWorker = new SwingWorker<List<Object[]>, Void>() {
@Override
protected List<String> doInBackground() {
protected List<Object[]> doInBackground() {
ConcurrentHashMap<String, Map<String, List<String>>> dataMap = Config.globalDataMap;
return exportData(selectedHost, exportDir, dataMap);
}

@Override
protected void done() {
try {
List<String> taskStatusList = get();
List<Object[]> taskStatusList = get();
if (!taskStatusList.isEmpty()) {
String exportStatusMessage = String.format("Exported File List Status:\n%s", String.join("\n", taskStatusList));

JOptionPane.showMessageDialog(Databoard.this, generateTaskStatusPane(exportStatusMessage), "Info", JOptionPane.INFORMATION_MESSAGE);
JOptionPane.showMessageDialog(Databoard.this, generateTaskStatusPane(taskStatusList), "Info", JOptionPane.INFORMATION_MESSAGE);
}
} catch (Exception ignored) {
}
Expand All @@ -409,17 +409,36 @@ protected void done() {
exportActionWorker.execute();
}

private JScrollPane generateTaskStatusPane(String message) {
JTextArea textArea = new JTextArea(message);
textArea.setEditable(false);
textArea.setLineWrap(true);
JScrollPane scrollPane = new JScrollPane(textArea);
scrollPane.setPreferredSize(new Dimension(400, 200));
private JScrollPane generateTaskStatusPane(List<Object[]> dataList) {
String[] columnNames = {"#", "Filename", "Status"};
DefaultTableModel taskStatusTableModel = new DefaultTableModel(columnNames, 0);
JTable taskStatusTable = new JTable(taskStatusTableModel);

for (Object[] data : dataList) {
int rowCount = taskStatusTable.getRowCount();
int id = rowCount > 0 ? (Integer) taskStatusTable.getValueAt(rowCount - 1, 0) + 1 : 1;
Object[] rowData = new Object[data.length + 1];
rowData[0] = id;
System.arraycopy(data, 0, rowData, 1, data.length);
taskStatusTableModel.addRow(rowData);
}

TableRowSorter<DefaultTableModel> sorter = new TableRowSorter<>(taskStatusTableModel);
taskStatusTable.setRowSorter(sorter);

JScrollPane scrollPane = new JScrollPane(taskStatusTable);
scrollPane.setBorder(new TitledBorder("Task status"));
scrollPane.setPreferredSize(new Dimension(500, 300));

int paneWidth = scrollPane.getPreferredSize().width;
taskStatusTable.getColumnModel().getColumn(0).setPreferredWidth((int) (paneWidth * 0.1));
taskStatusTable.getColumnModel().getColumn(1).setPreferredWidth((int) (paneWidth * 0.7));
taskStatusTable.getColumnModel().getColumn(2).setPreferredWidth((int) (paneWidth * 0.2));

return scrollPane;
}

private List<String> exportData(String selectedHost, String exportDir, Map<String, Map<String, List<String>>> dataMap) {
private List<Object[]> exportData(String selectedHost, String exportDir, Map<String, Map<String, List<String>>> dataMap) {
return dataMap.entrySet().stream()
.filter(entry -> selectedHost.equals("*") || StringProcessor.matchesHostPattern(entry.getKey(), selectedHost))
.filter(entry -> !entry.getKey().contains("*"))
Expand All @@ -428,7 +447,7 @@ private List<String> exportData(String selectedHost, String exportDir, Map<Strin
.collect(Collectors.toList());
}

private String exportEntry(Map.Entry<String, Map<String, List<String>>> entry, String exportDir) {
private Object[] exportEntry(Map.Entry<String, Map<String, List<String>>> entry, String exportDir) {
String key = entry.getKey();
Map<String, List<String>> ruleMap = entry.getValue();

Expand All @@ -442,7 +461,7 @@ private String exportEntry(Map.Entry<String, Map<String, List<String>>> entry, S
.collect(Collectors.toMap(
messageEntry -> messageEntry,
messageEntry -> StringProcessor.getRandomUUID(),
(existing, replacement) -> existing // 在冲突时保留现有的映射
(existing, replacement) -> existing
));

Map<String, Map<String, Object>> httpMap = processEntries(
Expand All @@ -463,7 +482,7 @@ private String exportEntry(Map.Entry<String, Map<String, List<String>>> entry, S
String filename = String.format("%s/%s-%s.hae", exportDir, StringProcessor.getCurrentTime(), hostName);
boolean createdStatus = projectProcessor.createHaeFile(filename, key, ruleMap, urlMap, httpMap);

return String.format("Filename: %s, Status: %s", filename, createdStatus);
return new Object[]{filename, createdStatus};
}


Expand Down Expand Up @@ -507,9 +526,9 @@ private void importActionPerformed(ActionEvent e) {
importActionWorker.cancel(true);
}

importActionWorker = new SwingWorker<List<String>, Void>() {
importActionWorker = new SwingWorker<List<Object[]>, Void>() {
@Override
protected List<String> doInBackground() {
protected List<Object[]> doInBackground() {
List<String> filesWithExtension = findFilesWithExtension(new File(exportDir), ".hae");
return filesWithExtension.stream()
.map(Databoard.this::importData)
Expand All @@ -519,10 +538,9 @@ protected List<String> doInBackground() {
@Override
protected void done() {
try {
List<String> taskStatusList = get();
List<Object[]> taskStatusList = get();
if (!taskStatusList.isEmpty()) {
String importStatusMessage = "Imported File List Status:\n" + String.join("\n", taskStatusList);
JOptionPane.showMessageDialog(Databoard.this, generateTaskStatusPane(importStatusMessage), "Info", JOptionPane.INFORMATION_MESSAGE);
JOptionPane.showMessageDialog(Databoard.this, generateTaskStatusPane(taskStatusList), "Info", JOptionPane.INFORMATION_MESSAGE);
}
} catch (Exception ignored) {
}
Expand All @@ -532,7 +550,7 @@ protected void done() {
importActionWorker.execute();
}

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

HaeFileContent haeFileContent = projectProcessor.readHaeFile(filename);
Expand Down Expand Up @@ -568,7 +586,7 @@ private String importData(String filename) {
}
}

return String.format("Filename: %s, Status: %s", filename, readStatus);
return new Object[]{filename, readStatus};
}

private List<String> findFilesWithExtension(File directory, String extension) {
Expand Down Expand Up @@ -648,6 +666,8 @@ private void clearActionPerformed(ActionEvent e) {
}

messageTableModel.deleteByHost(host);

hostTextField.setText("");
}
}
}
27 changes: 19 additions & 8 deletions src/main/java/hae/component/board/message/MessageTableModel.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -435,7 +435,7 @@ public String getColumnName(int columnIndex) {

public class MessageTable extends JTable {
private MessageEntry messageEntry;
private SwingWorker<Object, Void> currentWorker;
private SwingWorker<ByteArray[], Void> currentWorker;
private int lastSelectedIndex = -1;
private final HttpRequestEditor requestEditor;
private final HttpResponseEditor responseEditor;
Expand All @@ -450,16 +450,13 @@ public MessageTable(TableModel messageTableModel, HttpRequestEditor requestEdito
public void changeSelection(int row, int col, boolean toggle, boolean extend) {
super.changeSelection(row, col, toggle, extend);

requestEditor.setRequest(HttpRequest.httpRequest("Loading..."));
responseEditor.setResponse(HttpResponse.httpResponse("Loading..."));

if (currentWorker != null && !currentWorker.isDone()) {
currentWorker.cancel(true);
}

currentWorker = new SwingWorker<>() {
@Override
protected Void doInBackground() {
protected ByteArray[] doInBackground() {
int selectedIndex = convertRowIndexToModel(row);
if (lastSelectedIndex != selectedIndex) {
lastSelectedIndex = selectedIndex;
Expand All @@ -470,14 +467,28 @@ protected Void doInBackground() {
ByteArray requestByte = httpRequestResponse.request().toByteArray();
ByteArray responseByte = httpRequestResponse.response().toByteArray();

requestEditor.setRequest(HttpRequest.httpRequest(messageEntry.getRequestResponse().httpService(), requestByte));
responseEditor.setResponse(HttpResponse.httpResponse(responseByte));

ByteArray[] httpByteArray = new ByteArray[2];
httpByteArray[0] = requestByte;
httpByteArray[1] = responseByte;
return httpByteArray;
}

return null;
}

@Override
protected void done() {
try {
ByteArray[] retByteArray = get();
if (retByteArray != null) {
requestEditor.setRequest(HttpRequest.httpRequest(messageEntry.getRequestResponse().httpService(), retByteArray[0]));
responseEditor.setResponse(HttpResponse.httpResponse(retByteArray[1]));
}
} catch (Exception ignored) {
}
}
};

currentWorker.execute();
}
}
Expand Down
2 changes: 1 addition & 1 deletion src/main/java/hae/component/board/table/AIPower.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ public class AIPower {
public AIPower(MontoyaApi api, ConfigLoader configLoader, String aiModel, String aiBaseUrl, String[] apiKey) {
this.api = api;
this.configLoader = configLoader;
this.httpUtils = new HttpUtils(api);
this.httpUtils = new HttpUtils(api, configLoader);
this.aiModel = aiModel;
this.aiBaseUrl = aiBaseUrl;

Expand Down
34 changes: 24 additions & 10 deletions src/main/java/hae/component/config/Config.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,24 +27,23 @@ public class Config extends JPanel {
private final MontoyaApi api;
private final ConfigLoader configLoader;
private final Rules rules;
private JTextField addTextField;
private final String defaultText = "Enter a new item";
private final GridBagConstraints constraints = new GridBagConstraints();

public Config(MontoyaApi api, ConfigLoader configLoader, Rules rules) {
this.api = api;
this.configLoader = configLoader;
this.rules = rules;

constraints.weightx = 1.0;
constraints.fill = GridBagConstraints.HORIZONTAL;

initComponents();
}

private void initComponents() {
setLayout(new BorderLayout());

GridBagConstraints constraints = new GridBagConstraints();
constraints.weightx = 1.0;
constraints.fill = GridBagConstraints.HORIZONTAL;

JPanel ruleInfoPanel = new JPanel(new GridBagLayout());
ruleInfoPanel.setBorder(new EmptyBorder(10, 15, 5, 15));

Expand All @@ -67,7 +66,7 @@ private void initComponents() {
constraints.gridx = 1;
JTabbedPane configTabbedPanel = new JTabbedPane();

String[] settingMode = new String[]{"Exclude suffix", "Block host"};
String[] settingMode = new String[]{"Exclude suffix", "Block host", "Exclude status"};
JPanel settingPanel = createConfigTablePanel(settingMode, "Setting");
JPanel scopePanel = getScopePanel();
JScrollPane scopeScrollPane = new JScrollPane(scopePanel);
Expand Down Expand Up @@ -148,6 +147,12 @@ public void tableChanged(TableModelEvent e) {
configLoader.setBlockHost(values);
}
}

if (selected.equals("Exclude status")) {
if (!values.equals(configLoader.getExcludeStatus()) && !values.isEmpty()) {
configLoader.setExcludeStatus(values);
}
}
}
};
}
Expand All @@ -166,6 +171,10 @@ public void actionPerformed(ActionEvent e) {
if (selected.equals("Block host")) {
addDataToTable(configLoader.getBlockHost().replaceAll("\\|", "\r\n"), model);
}

if (selected.equals("Exclude status")) {
addDataToTable(configLoader.getExcludeStatus().replaceAll("\\|", "\r\n"), model);
}
}
};
}
Expand Down Expand Up @@ -211,6 +220,10 @@ public void actionPerformed(ActionEvent e) {
}

private JPanel createConfigTablePanel(String[] mode, String type) {
GridBagConstraints constraints = new GridBagConstraints();
constraints.weightx = 1.0;
constraints.fill = GridBagConstraints.HORIZONTAL;

JPanel settingPanel = new JPanel(new BorderLayout());
DefaultTableModel model = new DefaultTableModel();

Expand Down Expand Up @@ -255,7 +268,7 @@ private JPanel createConfigTablePanel(String[] mode, String type) {
constraints.gridy = 4;
buttonPanel.add(clearButton, constraints);

addTextField = new JTextField();
JTextField addTextField = new JTextField();
UIEnhancer.setTextFieldPlaceholder(addTextField, defaultText);

inputPanelB.add(addTextField, BorderLayout.CENTER);
Expand All @@ -266,13 +279,13 @@ private JPanel createConfigTablePanel(String[] mode, String type) {
settingPanel.add(inputPanel, BorderLayout.CENTER);


addButton.addActionListener(e -> addActionPerformed(e, model));
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);
addActionPerformed(null, model, addTextField);
}
}
});
Expand Down Expand Up @@ -372,8 +385,9 @@ public void updateScope(JCheckBox checkBox) {
configLoader.setScope(String.join("|", HaEScope));
}

private void addActionPerformed(ActionEvent e, DefaultTableModel model) {
private void addActionPerformed(ActionEvent e, DefaultTableModel model, JTextField addTextField) {
String addTextFieldText = addTextField.getText();
api.logging().logToOutput(addTextFieldText);
if (!addTextFieldText.equals(defaultText)) {
addDataToTable(addTextFieldText, model);
}
Expand Down
Loading

0 comments on commit 4f04013

Please sign in to comment.