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

Version: 4.0 Update

Browse Source
This commit is contained in:
gh0stkey
2024-12-21 15:34:45 +08:00
parent 1f1ca99f10
commit daacb2e146
24 changed files with 469 additions and 1427 deletions
Download Patch File Download Diff File Expand all files Collapse all files

110
src/main/java/hae/component/Config.java
View File

@@ -29,7 +29,6 @@ public class Config extends JPanel {
private final ConfigLoader configLoader;
private final MessageTableModel messageTableModel;
private final Rules rules;
private final String defaultText = "Enter a new item";
private Registration activeHandler;
private Registration passiveHandler;
@@ -61,22 +60,22 @@ public class Config extends JPanel {
pathTextField.setEditable(false);
pathTextField.setText(configLoader.getRulesFilePath());
JButton reloadButton = new JButton("Reload");
JButton updateButton = new JButton("Update");
JButton reinitButton = new JButton("Reinit");
ruleInfoPanel.add(ruleLabel);
ruleInfoPanel.add(pathTextField, constraints);
ruleInfoPanel.add(Box.createHorizontalStrut(5));
ruleInfoPanel.add(reloadButton);
ruleInfoPanel.add(reinitButton);
ruleInfoPanel.add(Box.createHorizontalStrut(5));
ruleInfoPanel.add(updateButton);
ruleInfoPanel.add(reloadButton);
reloadButton.addActionListener(this::reloadActionPerformed);
updateButton.addActionListener(this::onlineUpdateActionPerformed);
reinitButton.addActionListener(this::reinitActionPerformed);
constraints.gridx = 1;
JTabbedPane configTabbedPanel = new JTabbedPane();
String[] settingMode = new String[]{"Exclude suffix", "Block host", "Exclude status"};
JPanel settingPanel = createConfigTablePanel(settingMode, "Setting");
JPanel settingPanel = createConfigTablePanel(settingMode);
JPanel northPanel = new JPanel(new BorderLayout());
@@ -105,38 +104,6 @@ public class Config extends JPanel {
settingPanel.add(northPanel, BorderLayout.NORTH);
configTabbedPanel.add("Setting", settingPanel);
String[] aiMode = new String[]{"Alibaba", "Moonshot"};
JPanel aiPanel = createConfigTablePanel(aiMode, "AI+");
JTextArea promptTextArea = new JTextArea();
promptTextArea.setLineWrap(true);
promptTextArea.getDocument().addDocumentListener(new DocumentListener() {
@Override
public void insertUpdate(DocumentEvent e) {
onTextChange();
}
@Override
public void removeUpdate(DocumentEvent e) {
onTextChange();
}
@Override
public void changedUpdate(DocumentEvent e) {
onTextChange();
}
private void onTextChange() {
String promptText = promptTextArea.getText();
configLoader.setAIPrompt(promptText);
}
});
promptTextArea.setText(configLoader.getAIPrompt());
JScrollPane promptScrollPane = new JScrollPane(promptTextArea);
promptScrollPane.setBorder(new TitledBorder("Prompt"));
promptScrollPane.setPreferredSize(new Dimension(0, 100));
aiPanel.add(promptScrollPane, BorderLayout.NORTH);
configTabbedPanel.add("AI+", aiPanel);
add(ruleInfoPanel, BorderLayout.NORTH);
add(configTabbedPanel, BorderLayout.CENTER);
}
@@ -256,47 +223,8 @@ public class Config extends JPanel {
};
}
private TableModelListener craeteAITableModelListener(JComboBox<String> setTypeComboBox, DefaultTableModel model) {
return new TableModelListener() {
@Override
public void tableChanged(TableModelEvent e) {
String selected = (String) setTypeComboBox.getSelectedItem();
String values = getFirstColumnDataAsString(model);
if (selected.equals("Alibaba")) {
if (!values.equals(configLoader.getAlibabaAIAPIKey()) && !values.isEmpty()) {
configLoader.setAlibabaAIAPIKey(values);
}
}
if (selected.equals("Moonshot")) {
if (!values.equals(configLoader.getMoonshotAIAPIKey()) && !values.isEmpty()) {
configLoader.setMoonshotAIAPIKey(values);
}
}
}
};
}
private ActionListener createAIActionListener(JComboBox<String> setTypeComboBox, DefaultTableModel model) {
return new ActionListener() {
@Override
public void actionPerformed(ActionEvent e) {
String selected = (String) setTypeComboBox.getSelectedItem();
model.setRowCount(0);
if (selected.equals("Alibaba")) {
addDataToTable(configLoader.getAlibabaAIAPIKey().replaceAll("\\|", "\r\n"), model);
}
if (selected.equals("Moonshot")) {
addDataToTable(configLoader.getMoonshotAIAPIKey().replaceAll("\\|", "\r\n"), model);
}
}
};
}
private JPanel createConfigTablePanel(String[] mode, String type) {
private JPanel createConfigTablePanel(String[] mode) {
GridBagConstraints constraints = new GridBagConstraints();
constraints.weightx = 1.0;
constraints.fill = GridBagConstraints.HORIZONTAL;
@@ -327,9 +255,9 @@ public class Config extends JPanel {
JComboBox<String> setTypeComboBox = new JComboBox<>();
setTypeComboBox.setModel(new DefaultComboBoxModel<>(mode));
model.addTableModelListener(type.equals("AI+") ? craeteAITableModelListener(setTypeComboBox, model) : craeteSettingTableModelListener(setTypeComboBox, model));
model.addTableModelListener(craeteSettingTableModelListener(setTypeComboBox, model));
setTypeComboBox.addActionListener(type.equals("AI+") ? createAIActionListener(setTypeComboBox, model) : createSettingActionListener(setTypeComboBox, model));
setTypeComboBox.addActionListener(createSettingActionListener(setTypeComboBox, model));
setTypeComboBox.setSelectedItem(mode[0]);
@@ -346,6 +274,7 @@ public class Config extends JPanel {
buttonPanel.add(clearButton, constraints);
JTextField addTextField = new JTextField();
String defaultText = "Enter a new item";
UIEnhancer.setTextFieldPlaceholder(addTextField, defaultText);
inputPanelB.add(addTextField, BorderLayout.CENTER);
@@ -390,7 +319,7 @@ public class Config extends JPanel {
JPanel settingMainPanel = new JPanel(new BorderLayout());
settingMainPanel.setBorder(new EmptyBorder(5, 15, 10, 15));
JScrollPane settingScroller = new JScrollPane(settingPanel);
settingScroller.setBorder(new TitledBorder(type.equals("AI+") ? "API Key" : "Setting"));
settingScroller.setBorder(new TitledBorder("Setting"));
settingMainPanel.add(settingScroller, BorderLayout.CENTER);
return settingMainPanel;
@@ -492,16 +421,17 @@ public class Config extends JPanel {
}
}
private void onlineUpdateActionPerformed(ActionEvent e) {
// 添加提示框防止用户误触导致配置更新
int retCode = JOptionPane.showConfirmDialog(this, "Do you want to update rules?", "Info", JOptionPane.YES_NO_OPTION);
if (retCode == JOptionPane.YES_OPTION) {
configLoader.initRulesByNet();
reloadActionPerformed(null);
}
}
private void reloadActionPerformed(ActionEvent e) {
rules.reloadRuleGroup();
}
private void reinitActionPerformed(ActionEvent e) {
int retCode = JOptionPane.showConfirmDialog(this, "Do you want to reinitialize rules? This action will overwrite your existing rules.", "Info", JOptionPane.YES_NO_OPTION);
if (retCode == JOptionPane.YES_OPTION) {
boolean ret = configLoader.initRules();
if (ret) {
rules.reloadRuleGroup();
}
}
}
}

3
src/main/java/hae/component/Main.java
View File

@@ -81,7 +81,6 @@ public class Main extends JPanel {
ImageIcon originalIcon = new ImageIcon(imageURL);
Image originalImage = originalIcon.getImage();
Image scaledImage = originalImage.getScaledInstance(30, 20, Image.SCALE_FAST);
ImageIcon scaledIcon = new ImageIcon(scaledImage);
return scaledIcon;
return new ImageIcon(scaledImage);
}
}

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

@@ -2,49 +2,35 @@ package hae.component.board;
import burp.api.montoya.MontoyaApi;
import hae.Config;
import hae.component.board.message.MessageEntry;
import hae.component.board.message.MessageTableModel;
import hae.component.board.message.MessageTableModel.MessageTable;
import hae.component.board.table.Datatable;
import hae.instances.http.utils.RegularMatcher;
import hae.utils.ConfigLoader;
import hae.utils.UIEnhancer;
import hae.utils.project.ProjectProcessor;
import hae.utils.project.model.HaeFileContent;
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;
import java.awt.*;
import java.awt.event.*;
import java.io.File;
import java.util.List;
import java.util.*;
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;
public class Databoard extends JPanel {
private final MontoyaApi api;
private final ConfigLoader configLoader;
private final ProjectProcessor projectProcessor;
private final MessageTableModel messageTableModel;
private JTextField hostTextField;
private JTabbedPane dataTabbedPane;
private JSplitPane splitPane;
private MessageTable messageTable;
private JProgressBar progressBar;
private static Boolean isMatchHost = false;
private final DefaultComboBoxModel comboBoxModel = new DefaultComboBoxModel();
@@ -52,15 +38,10 @@ public class Databoard extends JPanel {
private SwingWorker<Map<String, List<String>>, Void> handleComboBoxWorker;
private SwingWorker<Void, Void> applyHostFilterWorker;
private SwingWorker<List<Object[]>, Void> exportActionWorker;
private SwingWorker<List<Object[]>, Void> importActionWorker;
private final String defaultText = "Please enter the host";
public Databoard(MontoyaApi api, ConfigLoader configLoader, MessageTableModel messageTableModel) {
this.api = api;
this.configLoader = configLoader;
this.projectProcessor = new ProjectProcessor(api);
this.messageTableModel = messageTableModel;
initComponents();
@@ -69,25 +50,22 @@ public class Databoard extends JPanel {
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, 0};
((GridBagLayout) getLayout()).rowHeights = new int[]{0, 65, 20, 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, 0.0, 1.0E-4};
((GridBagLayout) getLayout()).rowWeights = new double[]{0.0, 1.0, 0.0, 1.0E-4};
JLabel hostLabel = new JLabel("Host:");
JButton clearButton = new JButton("Clear");
JButton exportButton = new JButton("Export");
JButton importButton = new JButton("Import");
JButton actionButton = new JButton("Action");
JPanel menuPanel = new JPanel(new GridLayout(3, 1, 0, 5));
JPanel menuPanel = new JPanel(new GridLayout(1, 1, 0, 5));
menuPanel.setBorder(BorderFactory.createEmptyBorder(3, 3, 3, 3));
JPopupMenu menu = new JPopupMenu();
menuPanel.add(clearButton);
menuPanel.add(exportButton);
menuPanel.add(importButton);
menu.add(menuPanel);
hostTextField = new JTextField();
String defaultText = "Please enter the host";
UIEnhancer.setTextFieldPlaceholder(hostTextField, defaultText);
splitPane = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT);
@@ -102,10 +80,7 @@ public class Databoard extends JPanel {
});
clearButton.addActionListener(this::clearActionPerformed);
exportButton.addActionListener(this::exportActionPerformed);
importButton.addActionListener(this::importActionPerformed);
progressBar = new JProgressBar();
splitPane.addComponentListener(new ComponentAdapter() {
@Override
@@ -115,7 +90,6 @@ 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));
@@ -124,12 +98,9 @@ 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, 1, 0.0, 1.0,
add(splitPane, new GridBagConstraints(1, 1, 3, 2, 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));
@@ -149,24 +120,6 @@ public class Databoard extends JPanel {
columnModel.getColumn(5).setPreferredWidth((int) (totalWidth * 0.1));
}
private void setProgressBar(boolean status) {
setProgressBar(status, progressBar, "Loading ...");
}
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 setAutoMatch() {
hostComboBox.setSelectedItem(null);
hostComboBox.addActionListener(this::handleComboBoxAction);
@@ -202,8 +155,6 @@ 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()) {
@@ -357,11 +308,6 @@ public class Databoard extends JPanel {
return null;
}
@Override
protected void done() {
setProgressBar(false);
}
};
applyHostFilterWorker.execute();
@@ -374,263 +320,6 @@ public class Databoard extends JPanel {
return new ArrayList<>();
}
private void exportActionPerformed(ActionEvent e) {
String selectedHost = hostTextField.getText().trim();
if (selectedHost.isEmpty()) {
return;
}
String exportDir = selectDirectory(true);
if (exportDir.isEmpty()) {
return;
}
if (exportActionWorker != null && !exportActionWorker.isDone()) {
exportActionWorker.cancel(true);
}
exportActionWorker = new SwingWorker<List<Object[]>, Void>() {
@Override
protected List<Object[]> doInBackground() {
ConcurrentHashMap<String, Map<String, List<String>>> dataMap = Config.globalDataMap;
return exportData(selectedHost, exportDir, dataMap);
}
@Override
protected void done() {
try {
List<Object[]> taskStatusList = get();
if (!taskStatusList.isEmpty()) {
JOptionPane.showMessageDialog(Databoard.this, generateTaskStatusPane(taskStatusList), "Info", JOptionPane.INFORMATION_MESSAGE);
}
} catch (Exception ignored) {
}
}
};
exportActionWorker.execute();
}
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 = taskStatusTableModel.getRowCount();
int id = rowCount > 0 ? (Integer) taskStatusTableModel.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<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("*"))
.map(entry -> exportEntry(entry, exportDir))
.filter(Objects::nonNull)
.collect(Collectors.toList());
}
private Object[] exportEntry(Map.Entry<String, Map<String, List<String>>> entry, String exportDir) {
String key = entry.getKey();
Map<String, List<String>> ruleMap = entry.getValue();
if (ruleMap == null || ruleMap.isEmpty()) {
return null;
}
List<MessageEntry> messageEntryList = messageTableModel.getLogs();
Map<MessageEntry, String> entryUUIDMap = messageEntryList.stream()
.collect(Collectors.toMap(
messageEntry -> messageEntry,
messageEntry -> StringProcessor.getRandomUUID(),
(existing, replacement) -> existing
));
Map<String, Map<String, Object>> httpMap = processEntries(
messageEntryList,
key,
entryUUIDMap,
this::createHttpItemMap
);
Map<String, Map<String, Object>> urlMap = processEntries(
messageEntryList,
key,
entryUUIDMap,
this::creteUrlItemMap
);
String hostName = key.replace(":", "_");
String filename = String.format("%s/%s-%s.hae", exportDir, StringProcessor.getCurrentTime(), hostName);
boolean createdStatus = projectProcessor.createHaeFile(filename, key, ruleMap, urlMap, httpMap);
return new Object[]{filename, createdStatus};
}
private Map<String, Map<String, Object>> processEntries(List<MessageEntry> messageEntryList, String key, Map<MessageEntry, String> entryUUIDMap, Function<MessageEntry, Map<String, Object>> mapFunction) {
return messageEntryList.stream()
.filter(messageEntry -> !StringProcessor.getHostByUrl(messageEntry.getUrl()).isEmpty())
.filter(messageEntry -> StringProcessor.getHostByUrl(messageEntry.getUrl()).equals(key))
.collect(Collectors.toMap(
entryUUIDMap::get,
mapFunction,
(existing, replacement) -> existing
));
}
private Map<String, Object> creteUrlItemMap(MessageEntry entry) {
Map<String, Object> urlItemMap = new LinkedHashMap<>();
urlItemMap.put("url", entry.getUrl());
urlItemMap.put("method", entry.getMethod());
urlItemMap.put("status", entry.getStatus());
urlItemMap.put("length", entry.getLength());
urlItemMap.put("comment", entry.getComment());
urlItemMap.put("color", entry.getColor());
urlItemMap.put("size", String.valueOf(entry.getRequestResponse().request().toByteArray().length()));
return urlItemMap;
}
private Map<String, Object> createHttpItemMap(MessageEntry entry) {
Map<String, Object> httpItemMap = new LinkedHashMap<>();
httpItemMap.put("request", entry.getRequestResponse().request().toByteArray().getBytes());
httpItemMap.put("response", entry.getRequestResponse().response().toByteArray().getBytes());
return httpItemMap;
}
private void importActionPerformed(ActionEvent e) {
String exportDir = selectDirectory(false);
if (exportDir.isEmpty()) {
return;
}
if (importActionWorker != null && !importActionWorker.isDone()) {
importActionWorker.cancel(true);
}
importActionWorker = new SwingWorker<List<Object[]>, Void>() {
@Override
protected List<Object[]> doInBackground() {
List<String> filesWithExtension = findFilesWithExtension(new File(exportDir), ".hae");
return filesWithExtension.stream()
.map(Databoard.this::importData)
.collect(Collectors.toList());
}
@Override
protected void done() {
try {
List<Object[]> taskStatusList = get();
if (!taskStatusList.isEmpty()) {
JOptionPane.showMessageDialog(Databoard.this, generateTaskStatusPane(taskStatusList), "Info", JOptionPane.INFORMATION_MESSAGE);
}
} catch (Exception ignored) {
}
}
};
importActionWorker.execute();
}
private Object[] importData(String filename) {
ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2);
HaeFileContent haeFileContent = projectProcessor.readHaeFile(filename);
boolean readStatus = haeFileContent != null;
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) -> {
tasks.add(() -> {
String url = urlItemMap.get("url");
String comment = urlItemMap.get("comment");
String color = urlItemMap.get("color");
String length = urlItemMap.get("length");
String method = urlItemMap.get("method");
String status = urlItemMap.get("status");
String path = haeFileContent.getHttpPath();
messageTableModel.add(null, url, method, status, length, comment, color, key, path);
return null;
});
});
executor.invokeAll(tasks);
} catch (Exception e) {
api.logging().logToError("importData: " + e.getMessage());
} finally {
executor.shutdown();
}
}
return new Object[]{filename, readStatus};
}
private List<String> findFilesWithExtension(File directory, String extension) {
List<String> filePaths = new ArrayList<>();
if (directory.isDirectory()) {
File[] files = directory.listFiles();
if (files != null) {
for (File file : files) {
if (file.isDirectory()) {
filePaths.addAll(findFilesWithExtension(file, extension));
} else if (file.isFile() && file.getName().toLowerCase().endsWith(extension)) {
filePaths.add(file.getAbsolutePath());
}
}
}
} else {
filePaths.add(directory.getAbsolutePath());
}
return filePaths;
}
private String selectDirectory(boolean forDirectories) {
JFileChooser chooser = new JFileChooser();
chooser.setCurrentDirectory(new java.io.File(configLoader.getRulesFilePath()));
chooser.setDialogTitle(String.format("Select a Directory%s", forDirectories ? "" : " or File"));
FileNameExtensionFilter filter = new FileNameExtensionFilter(".hae Files", "hae");
chooser.addChoosableFileFilter(filter);
chooser.setFileFilter(filter);
chooser.setFileSelectionMode(forDirectories ? JFileChooser.DIRECTORIES_ONLY : JFileChooser.FILES_AND_DIRECTORIES);
chooser.setAcceptAllFileFilterUsed(!forDirectories);
if (chooser.showOpenDialog(this) == JFileChooser.APPROVE_OPTION) {
File selectedDirectory = chooser.getSelectedFile();
return selectedDirectory.getAbsolutePath();
}
return "";
}
private void clearActionPerformed(ActionEvent e) {
int retCode = JOptionPane.showConfirmDialog(this, "Do you want to clear data?", "Info",
JOptionPane.YES_NO_OPTION);
@@ -638,7 +327,6 @@ public class Databoard extends JPanel {
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("*")) {

14
src/main/java/hae/component/board/message/MessageEntry.java
View File

@@ -11,10 +11,8 @@ public class MessageEntry {
private final String status;
private final String color;
private final String method;
private final String hash;
private final String path;
MessageEntry(HttpRequestResponse requestResponse, String method, String url, String comment, String length, String color, String status, String hash, String path) {
MessageEntry(HttpRequestResponse requestResponse, String method, String url, String comment, String length, String color, String status) {
this.requestResponse = requestResponse;
this.method = method;
this.url = url;
@@ -22,8 +20,6 @@ public class MessageEntry {
this.length = length;
this.color = color;
this.status = status;
this.hash = hash;
this.path = path;
}
public String getColor() {
@@ -53,12 +49,4 @@ public class MessageEntry {
public HttpRequestResponse getRequestResponse() {
return this.requestResponse;
}
public String getHash() {
return this.hash;
}
public String getPath() {
return this.path;
}
}

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

@@ -5,13 +5,14 @@ import burp.api.montoya.http.message.HttpHeader;
import burp.api.montoya.http.message.HttpRequestResponse;
import burp.api.montoya.http.message.requests.HttpRequest;
import burp.api.montoya.http.message.responses.HttpResponse;
import burp.api.montoya.persistence.PersistedObject;
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.utils.ConfigLoader;
import hae.utils.project.FileProcessor;
import hae.utils.DataManager;
import hae.utils.string.HashCalculator;
import hae.utils.string.StringProcessor;
@@ -97,10 +98,10 @@ 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, String hash, String path) {
public 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, hash, path);
MessageEntry logEntry = new MessageEntry(messageInfo, method, url, comment, length, color, status);
byte[] reqByteA = new byte[0];
byte[] resByteA = new byte[0];
@@ -134,6 +135,18 @@ public class MessageTableModel extends AbstractTableModel {
}
if (!isDuplicate) {
if (flag) {
DataManager dataManager = new DataManager(api);
// 数据存储在BurpSuite空间内
PersistedObject persistedObject = PersistedObject.persistedObject();
persistedObject.setHttpRequestResponse("messageInfo", messageInfo);
persistedObject.setString("comment", comment);
persistedObject.setString("color", color);
String uuidIndex = StringProcessor.getRandomUUID();
dataManager.putData("message", uuidIndex, persistedObject);
}
// 添加进日志
log.add(logEntry);
}
}
@@ -177,11 +190,10 @@ public class MessageTableModel extends AbstractTableModel {
filteredLog.clear();
log.forEach(entry -> {
MessageEntry finalEntry = getEntryByFile(entry);
String host = StringProcessor.getHostByUrl(finalEntry.getUrl());
String host = StringProcessor.getHostByUrl(entry.getUrl());
if (!host.isEmpty()) {
if (StringProcessor.matchesHostPattern(host, filterText) || filterText.contains("*")) {
filteredLog.add(finalEntry);
filteredLog.add(entry);
}
}
});
@@ -189,34 +201,6 @@ public class MessageTableModel extends AbstractTableModel {
fireTableDataChanged();
}
private MessageEntry getEntryByFile(MessageEntry entry) {
HttpRequestResponse requestResponse = entry.getRequestResponse();
if (requestResponse == null) {
String url = entry.getUrl();
String method = entry.getMethod();
String status = entry.getStatus();
String comment = entry.getComment();
String color = entry.getColor();
String path = entry.getPath();
String hash = entry.getHash();
int length = Integer.parseInt(entry.getLength());
byte[] contents = FileProcessor.readFileContent(path, hash);
if (contents.length > length) {
byte[] response = Arrays.copyOf(contents, length);
byte[] request = Arrays.copyOfRange(contents, length, contents.length);
requestResponse = StringProcessor.createHttpRequestResponse(url, request, response);
int index = log.indexOf(entry);
entry = new MessageEntry(requestResponse, method, url, comment, String.valueOf(length), color, status, "", "");
log.set(index, entry);
}
}
return entry;
}
public void applyMessageFilter(String tableName, String filterText) {
filteredLog.clear();
for (MessageEntry entry : log) {

157
src/main/java/hae/component/board/table/AIPower.java
View File

@@ -1,157 +0,0 @@
package hae.component.board.table;
import burp.api.montoya.MontoyaApi;
import burp.api.montoya.http.RequestOptions;
import burp.api.montoya.http.message.HttpRequestResponse;
import burp.api.montoya.http.message.requests.HttpRequest;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.reflect.TypeToken;
import hae.Config;
import hae.utils.ConfigLoader;
import hae.utils.http.HttpUtils;
import okhttp3.*;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.lang.reflect.Type;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
public class AIPower {
private final MontoyaApi api;
private final HttpUtils httpUtils;
private final ConfigLoader configLoader;
private final String apiAuth;
private final String aiModel;
private final String aiBaseUrl;
public AIPower(MontoyaApi api, ConfigLoader configLoader, String aiModel, String aiBaseUrl, String[] apiKey) {
this.api = api;
this.configLoader = configLoader;
this.httpUtils = new HttpUtils(api, configLoader);
this.aiModel = aiModel;
this.aiBaseUrl = aiBaseUrl;
this.apiAuth = String.format("Bearer %s", apiKey[new Random().nextInt(apiKey.length)]);
}
// Stream Response
public String chatWithAPI(String ruleName, String data) {
OkHttpClient httpClient = new OkHttpClient();
String fileId = uploadFileToAIService(ruleName, data);
Gson gson = new Gson();
if (fileId != null) {
String chatUrl = String.format("%s/chat/completions", aiBaseUrl);
String chatMessage = generateJsonData(configLoader.getAIPrompt(), fileId);
Request request = new Request.Builder()
.url(chatUrl)
.header("Authorization", apiAuth)
.post(RequestBody.create(MediaType.parse("application/json"), chatMessage))
.build();
try (Response response = httpClient.newCall(request).execute()) {
if (!response.isSuccessful()) {
throw new IOException("Unexpected code " + response);
}
BufferedReader reader = new BufferedReader(new InputStreamReader(response.body().byteStream()));
StringBuilder chatReturn = new StringBuilder();
String line;
while ((line = reader.readLine()) != null) {
if (line.startsWith("data: ") && !line.contains("[DONE]")) {
String jsonData = line.substring(6);
Type type = new TypeToken<Map<String, Object>>() {
}.getType();
Map<String, Object> map = gson.fromJson(jsonData, type);
String content = getDeltaContent(map);
if (content != null) {
chatReturn.append(content);
}
}
}
deleteFileOnAIService(fileId);
return chatReturn.toString();
} catch (Exception e) {
return "";
}
}
return "";
}
private String getDeltaContent(Map<String, Object> map) {
List<Map<String, Map<String, String>>> choices = (List<Map<String, Map<String, String>>>) map.get("choices");
if (choices != null && !choices.isEmpty()) {
Map<String, String> delta = choices.get(0).get("delta");
return delta.get("content");
}
return null;
}
private String uploadFileToAIService(String ruleName, String data) {
String uploadUrl = String.format("%s/files", aiBaseUrl);
String uploadParam = "file";
String filename = "hae.txt";
String content = String.format(Config.userTextFormat, ruleName, data);
HttpRequest uploadFileRequest = httpUtils.generateRequestByMultipartUploadMethod(uploadUrl, uploadParam, filename, content).withAddedHeader("Authorization", apiAuth);
HttpRequestResponse uploadFileRequestResponse = api.http().sendRequest(uploadFileRequest, RequestOptions.requestOptions().withUpstreamTLSVerification());
String responseBody = uploadFileRequestResponse.response().bodyToString();
Pattern pattern = Pattern.compile("\"id\":\"(.*?)\",");
Matcher matcher = pattern.matcher(responseBody);
return matcher.find() ? matcher.group(1) : null;
}
private void deleteFileOnAIService(String fileId) {
String deleteFileUrl = String.format("%s/files/%s", aiBaseUrl, fileId);
HttpRequest deleteFileRequest = httpUtils.generateRequestByDeleteMethod(deleteFileUrl).withAddedHeader("Authorization", apiAuth);
api.http().sendRequest(deleteFileRequest, RequestOptions.requestOptions().withUpstreamTLSVerification());
}
private String getFileContentOnAiService(String fileId) {
String getFileContentUrl = String.format("%s/files/%s/content", aiBaseUrl, fileId);
HttpRequest getFileContentRequest = HttpRequest.httpRequestFromUrl(getFileContentUrl).withAddedHeader("Authorization", apiAuth);
HttpRequestResponse getFileRequestResponse = api.http().sendRequest(getFileContentRequest, RequestOptions.requestOptions().withUpstreamTLSVerification());
String responseBody = getFileRequestResponse.response().bodyToString();
Pattern pattern = Pattern.compile("\"content\":\"(.*?)\",\"file_type\"");
Matcher matcher = pattern.matcher(responseBody);
return matcher.find() ? matcher.group(1) : null;
}
private String generateJsonData(String prompt, String fileId) {
Map<String, Object> data = new HashMap<>();
data.put("model", aiModel);
data.put("stream", true);
data.put("messages", new Object[]{
new HashMap<String, Object>() {{
put("role", "system");
put("content", prompt);
}},
new HashMap<String, Object>() {{
put("role", "system");
put("content", aiModel.equals("qwen-long") ? String.format("fileid://%s", fileId) : getFileContentOnAiService(fileId));
}},
new HashMap<String, Object>() {{
put("role", "user");
put("content", "Start");
}}
});
Gson gson = new GsonBuilder().setPrettyPrinting().create();
return gson.toJson(data);
}
}

111
src/main/java/hae/component/board/table/Datatable.java
View File

@@ -1,9 +1,6 @@
package hae.component.board.table;
import burp.api.montoya.MontoyaApi;
import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
import hae.component.board.Databoard;
import hae.component.board.message.MessageTableModel;
import hae.utils.ConfigLoader;
import hae.utils.UIEnhancer;
@@ -11,22 +8,17 @@ import hae.utils.UIEnhancer;
import javax.swing.*;
import javax.swing.event.DocumentEvent;
import javax.swing.event.DocumentListener;
import javax.swing.event.PopupMenuEvent;
import javax.swing.event.PopupMenuListener;
import javax.swing.table.DefaultTableModel;
import javax.swing.table.TableColumn;
import javax.swing.table.TableRowSorter;
import java.awt.*;
import java.awt.datatransfer.Clipboard;
import java.awt.datatransfer.StringSelection;
import java.awt.event.ActionEvent;
import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import java.lang.reflect.Type;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.regex.Pattern;
public class Datatable extends JPanel {
@@ -39,15 +31,12 @@ public class Datatable extends JPanel {
private final TableRowSorter<DefaultTableModel> sorter;
private final JCheckBox searchMode = new JCheckBox("Reverse search");
private final String tabName;
private final JProgressBar progressBar;
private final JPopupMenu aiEmpoweredMenu;
private final JPanel footerPanel;
public Datatable(MontoyaApi api, ConfigLoader configLoader, String tabName, List<String> dataList) {
this.api = api;
this.configLoader = configLoader;
this.tabName = tabName;
this.progressBar = new JProgressBar();
String[] columnNames = {"#", "Information"};
this.dataTableModel = new DefaultTableModel(columnNames, 0);
@@ -56,15 +45,12 @@ public class Datatable extends JPanel {
this.sorter = new TableRowSorter<>(dataTableModel);
this.searchField = new JTextField(10);
this.secondSearchField = new JTextField(10);
this.aiEmpoweredMenu = new JPopupMenu();
this.footerPanel = new JPanel(new BorderLayout(0, 5));
initComponents(dataList);
}
private void initComponents(List<String> dataList) {
progressBar.setVisible(false);
// 设置ID排序
sorter.setComparator(0, new Comparator<Integer>() {
@Override
@@ -142,57 +128,17 @@ public class Datatable extends JPanel {
JButton settingsButton = new JButton("Settings");
setMenuShow(settingMenu, settingsButton);
// AI Empowered按钮
JPanel aiEmpoweredPanel = new JPanel(new GridLayout(2, 1));
aiEmpoweredPanel.setBorder(BorderFactory.createEmptyBorder(3, 3, 3, 3));
JButton empoweredByAlibabaButton = new JButton("Alibaba - QwenLong");
empoweredByAlibabaButton.addActionListener(e -> {
aiEmpoweredByAlibabaActionPerformed(e, tabName, getTableData(dataTable));
});
JButton empoweredByMoonshotButton = new JButton("Moonshot - Kimi");
empoweredByMoonshotButton.addActionListener(e -> {
aiEmpoweredByMoonshotActionPerformed(e, tabName, getTableData(dataTable));
});
aiEmpoweredPanel.add(empoweredByAlibabaButton);
aiEmpoweredPanel.add(empoweredByMoonshotButton);
aiEmpoweredMenu.add(aiEmpoweredPanel);
JButton aiEmpoweredButton = new JButton("AI Empowered");
setMenuShow(aiEmpoweredMenu, aiEmpoweredButton);
aiEmpoweredMenu.addPopupMenuListener(new PopupMenuListener() {
@Override
public void popupMenuWillBecomeVisible(PopupMenuEvent e) {
empoweredByAlibabaButton.setEnabled(!configLoader.getAlibabaAIAPIKey().isEmpty());
empoweredByMoonshotButton.setEnabled(!configLoader.getMoonshotAIAPIKey().isEmpty());
}
@Override
public void popupMenuWillBecomeInvisible(PopupMenuEvent e) {
}
@Override
public void popupMenuCanceled(PopupMenuEvent e) {
}
});
optionsPanel.add(settingsButton);
optionsPanel.add(Box.createHorizontalStrut(5));
optionsPanel.add(searchField);
optionsPanel.add(Box.createHorizontalStrut(5));
optionsPanel.add(secondSearchField);
optionsPanel.add(Box.createHorizontalStrut(5));
optionsPanel.add(aiEmpoweredButton);
footerPanel.setBorder(BorderFactory.createEmptyBorder(2, 3, 5, 3));
footerPanel.add(optionsPanel, BorderLayout.CENTER);
footerPanel.add(progressBar, BorderLayout.SOUTH);
add(scrollPane, BorderLayout.CENTER);
add(footerPanel, BorderLayout.SOUTH);
setProgressBar(false);
}
private void setMenuShow(JPopupMenu menu, JButton button) {
@@ -205,9 +151,6 @@ public class Datatable extends JPanel {
});
}
private void setProgressBar(boolean status) {
Databoard.setProgressBar(status, progressBar, "AI+ ...");
}
private void addRowToTable(Object[] data) {
int rowCount = dataTableModel.getRowCount();
@@ -218,60 +161,6 @@ public class Datatable extends JPanel {
dataTableModel.addRow(rowData);
}
private void aiEmpoweredByAlibabaActionPerformed(ActionEvent e, String ruleName, String data) {
AIPower aiPower = new AIPower(api, configLoader, "qwen-long", "https://dashscope.aliyuncs.com/compatible-mode/v1", configLoader.getAlibabaAIAPIKey().split("\\|"));
aiEmpoweredButtonAction(ruleName, data, aiPower);
}
private void aiEmpoweredByMoonshotActionPerformed(ActionEvent e, String ruleName, String data) {
AIPower aiPower = new AIPower(api, configLoader, "moonshot-v1-128k", "https://api.moonshot.cn/v1", configLoader.getMoonshotAIAPIKey().split("\\|"));
aiEmpoweredButtonAction(ruleName, data, aiPower);
}
private void aiEmpoweredButtonAction(String ruleName, String data, AIPower aiPower) {
progressBar.setVisible(true);
aiEmpoweredMenu.setVisible(true);
setProgressBar(true);
SwingWorker<String, Void> worker = new SwingWorker<String, Void>() {
@Override
protected String doInBackground() throws Exception {
return aiPower.chatWithAPI(ruleName, data);
}
@Override
protected void done() {
setProgressBar(false);
try {
String chatReturn = get();
if (!chatReturn.isEmpty()) {
Gson gson = new Gson();
Type type = new TypeToken<Map<String, Object>>() {
}.getType();
Map<String, List<String>> map = gson.fromJson(chatReturn, type);
dataTableModel.setRowCount(0);
for (String item : map.get("data")) {
if (!item.isEmpty()) {
addRowToTable(new Object[]{item});
}
}
JOptionPane.showMessageDialog(Datatable.this, "AI+ has completed the AI empowered work.", "AI+ Info", JOptionPane.INFORMATION_MESSAGE);
} else {
JOptionPane.showMessageDialog(Datatable.this, "AI+ returns null, please check!", "AI+ Info", JOptionPane.WARNING_MESSAGE);
}
} catch (Exception ignored) {
JOptionPane.showMessageDialog(Datatable.this, "AI+ returns error, please check!", "AI+ Info", JOptionPane.ERROR_MESSAGE);
}
}
};
worker.execute();
aiEmpoweredMenu.setVisible(false);
}
private void performSearch() {
RowFilter<Object, Object> firstRowFilter = applyFirstSearchFilter();
RowFilter<Object, Object> secondRowFilter = applySecondFilter();