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

Version: 3.3 Update

Browse Source
This commit is contained in:
gh0stkey
2024-07-23 09:21:30 +08:00
parent e4b7f86a0c
commit 5f54d1f461
19 changed files with 1166 additions and 644 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -37,26 +37,17 @@ public class Main extends JPanel {
// 新增Logo
JTabbedPane HaETabbedPane = new JTabbedPane();
HaETabbedPane.addTab("", getImageIcon(false), mainTabbedPane);
boolean isDarkBg = isDarkBg(HaETabbedPane);
HaETabbedPane.addTab("", getImageIcon(isDarkBg), mainTabbedPane);
// 中文Slogan：赋能白帽，高效作战
HaETabbedPane.addTab(" Highlighter and Extractor - Empower ethical hacker for efficient operations. ", null);
HaETabbedPane.setEnabledAt(1, false);
HaETabbedPane.addPropertyChangeListener("background", new PropertyChangeListener() {
@Override
public void propertyChange(PropertyChangeEvent e) {
boolean isDarkBg = isDarkBg();
boolean isDarkBg = isDarkBg(HaETabbedPane);
HaETabbedPane.setIconAt(0, getImageIcon(isDarkBg));
}
private boolean isDarkBg() {
Color bg = HaETabbedPane.getBackground();
int r = bg.getRed();
int g = bg.getGreen();
int b = bg.getBlue();
int avg = (r + g + b) / 3;
return avg < 128;
}
});
add(HaETabbedPane, new GridBagConstraints(0, 0, 1, 1, 0.0, 0.0,
@@ -70,6 +61,16 @@ public class Main extends JPanel {
mainTabbedPane.addTab("Config", new Config(api, configLoader, rules));
}
private boolean isDarkBg(JTabbedPane HaETabbedPane) {
Color bg = HaETabbedPane.getBackground();
int r = bg.getRed();
int g = bg.getGreen();
int b = bg.getBlue();
int avg = (r + g + b) / 3;
return avg < 128;
}
private ImageIcon getImageIcon(boolean isDark) {
ClassLoader classLoader = getClass().getClassLoader();
URL imageURL;

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

@@ -5,6 +5,7 @@ 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.project.ProjectProcessor;
@@ -46,7 +47,7 @@ public class Databoard extends JPanel {
private final DefaultComboBoxModel comboBoxModel = new DefaultComboBoxModel();
private final JComboBox hostComboBox = new JComboBox(comboBoxModel);
private SwingWorker<Boolean, Void> handleComboBoxWorker;
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;
@@ -63,7 +64,7 @@ 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, 25, 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, 0.0, 1.0E-4};
@@ -120,7 +121,7 @@ public class Databoard extends JPanel {
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, 3, 3, 1, 1.0, 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);
@@ -143,6 +144,11 @@ public class Databoard extends JPanel {
}
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);
@@ -150,7 +156,7 @@ public class Databoard extends JPanel {
progressBar.setStringPainted(true);
progressBar.setValue(progressBar.getMaximum());
} else {
progressBar.setString("Loading...");
progressBar.setString(showString);
progressBar.setStringPainted(true);
}
}
@@ -198,18 +204,27 @@ public class Databoard extends JPanel {
handleComboBoxWorker.cancel(true);
}
handleComboBoxWorker = new SwingWorker<Boolean, Void>() {
handleComboBoxWorker = new SwingWorker<Map<String, List<String>>, Void>() {
@Override
protected Boolean doInBackground() {
return populateTabbedPaneByHost(selectedHost);
protected Map<String, List<String>> doInBackground() {
return getSelectedMapByHost(selectedHost);
}
@Override
protected void done() {
if (!isCancelled()) {
try {
boolean status = get();
if (status) {
Map<String, List<String>> selectedDataMap = get();
if (!selectedDataMap.isEmpty()) {
dataTabbedPane.removeAll();
for (Map.Entry<String, List<String>> entry : selectedDataMap.entrySet()) {
String tabTitle = String.format("%s (%s)", entry.getKey(), entry.getValue().size());
Datatable datatablePanel = new Datatable(api, configLoader, entry.getKey(), entry.getValue());
datatablePanel.setTableListener(messageTableModel);
dataTabbedPane.addTab(tabTitle, datatablePanel);
}
JSplitPane messageSplitPane = messageTableModel.getSplitPane();
splitPane.setLeftComponent(dataTabbedPane);
splitPane.setRightComponent(messageSplitPane);
@@ -257,7 +272,7 @@ public class Databoard extends JPanel {
isMatchHost = false;
}
private boolean populateTabbedPaneByHost(String selectedHost) {
private Map<String, List<String>> getSelectedMapByHost(String selectedHost) {
ConcurrentHashMap<String, Map<String, List<String>>> dataMap = Config.globalDataMap;
Map<String, List<String>> selectedDataMap;
@@ -283,20 +298,7 @@ public class Databoard extends JPanel {
selectedDataMap = dataMap.get(selectedHost);
}
if (!selectedDataMap.isEmpty()) {
dataTabbedPane.removeAll();
for (Map.Entry<String, List<String>> entry : selectedDataMap.entrySet()) {
String tabTitle = String.format("%s (%s)", entry.getKey(), entry.getValue().size());
Datatable datatablePanel = new Datatable(api, entry.getKey(), entry.getValue());
datatablePanel.setTableListener(messageTableModel);
dataTabbedPane.addTab(tabTitle, datatablePanel);
}
return true;
}
return false;
return selectedDataMap;
}
private void filterComboBoxList() {
@@ -396,7 +398,8 @@ public class Databoard extends JPanel {
List<String> taskStatusList = get();
if (!taskStatusList.isEmpty()) {
String exportStatusMessage = String.format("Exported File List Status:\n%s", String.join("\n", taskStatusList));
JOptionPane.showConfirmDialog(Databoard.this, exportStatusMessage, "Info", JOptionPane.YES_OPTION);
JOptionPane.showMessageDialog(Databoard.this, generateTaskStatusPane(exportStatusMessage), "Info", JOptionPane.INFORMATION_MESSAGE);
}
} catch (Exception ignored) {
}
@@ -406,6 +409,16 @@ public class Databoard extends JPanel {
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));
return scrollPane;
}
private List<String> 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))
@@ -509,7 +522,7 @@ public class Databoard extends JPanel {
List<String> taskStatusList = get();
if (!taskStatusList.isEmpty()) {
String importStatusMessage = "Imported File List Status:\n" + String.join("\n", taskStatusList);
JOptionPane.showConfirmDialog(Databoard.this, importStatusMessage, "Info", JOptionPane.YES_OPTION);
JOptionPane.showMessageDialog(Databoard.this, generateTaskStatusPane(importStatusMessage), "Info", JOptionPane.INFORMATION_MESSAGE);
}
} catch (Exception ignored) {
}

207
src/main/java/hae/component/board/Datatable.java
View File

@@ -1,207 +0,0 @@
package hae.component.board;
import burp.api.montoya.MontoyaApi;
import hae.component.board.message.MessageTableModel;
import hae.utils.UIEnhancer;
import javax.swing.*;
import javax.swing.event.DocumentEvent;
import javax.swing.event.DocumentListener;
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.MouseAdapter;
import java.awt.event.MouseEvent;
import java.util.Comparator;
import java.util.List;
import java.util.regex.Pattern;
public class Datatable extends JPanel {
private final MontoyaApi api;
private final JTable dataTable;
private final DefaultTableModel dataTableModel;
private final JTextField searchField;
private final TableRowSorter<DefaultTableModel> sorter;
private final JCheckBox searchMode = new JCheckBox("Reverse search");
private final String tabName;
public Datatable(MontoyaApi api, String tabName, List<String> dataList) {
this.api = api;
this.tabName = tabName;
String[] columnNames = {"#", "Information"};
dataTableModel = new DefaultTableModel(columnNames, 0);
dataTable = new JTable(dataTableModel);
sorter = new TableRowSorter<>(dataTableModel);
searchField = new JTextField();
initComponents(dataList);
}
private void initComponents(List<String> dataList) {
// 设置ID排序
sorter.setComparator(0, new Comparator<Integer>() {
@Override
public int compare(Integer s1, Integer s2) {
return s1.compareTo(s2);
}
});
dataTable.setRowSorter(sorter);
TableColumn idColumn = dataTable.getColumnModel().getColumn(0);
idColumn.setMaxWidth(50);
for (String item : dataList) {
if (!item.isEmpty()) {
addRowToTable(new Object[]{item});
}
}
// 设置灰色默认文本
String searchText = "Search";
UIEnhancer.setTextFieldPlaceholder(searchField, searchText);
// 监听输入框内容输入、更新、删除
searchField.getDocument().addDocumentListener(new DocumentListener() {
@Override
public void insertUpdate(DocumentEvent e) {
performSearch();
}
@Override
public void removeUpdate(DocumentEvent e) {
performSearch();
}
@Override
public void changedUpdate(DocumentEvent e) {
performSearch();
}
});
// 设置布局
JScrollPane scrollPane = new JScrollPane(dataTable);
scrollPane.setVerticalScrollBarPolicy(ScrollPaneConstants.VERTICAL_SCROLLBAR_ALWAYS);
searchMode.addItemListener(e -> performSearch());
setLayout(new BorderLayout(0, 5));
JPanel optionsPanel = new JPanel();
optionsPanel.setBorder(BorderFactory.createEmptyBorder(2, 3, 5, 5));
optionsPanel.setLayout(new BoxLayout(optionsPanel, BoxLayout.X_AXIS));
// 新增复选框要在这修改rows
JPanel menuPanel = new JPanel(new GridLayout(1, 1));
menuPanel.setBorder(BorderFactory.createEmptyBorder(3, 3, 3, 3));
JPopupMenu menu = new JPopupMenu();
menuPanel.add(searchMode);
menu.add(menuPanel);
JButton settingsButton = new JButton("Settings");
settingsButton.addActionListener(e -> {
int x = settingsButton.getX();
int y = settingsButton.getY() - menu.getPreferredSize().height;
menu.show(settingsButton, x, y);
});
optionsPanel.add(settingsButton);
optionsPanel.add(Box.createHorizontalStrut(5));
optionsPanel.add(searchField);
add(scrollPane, BorderLayout.CENTER);
add(optionsPanel, BorderLayout.SOUTH);
}
private void addRowToTable(Object[] data) {
int rowCount = dataTableModel.getRowCount();
int id = rowCount > 0 ? (Integer) dataTableModel.getValueAt(rowCount - 1, 0) + 1 : 1;
Object[] rowData = new Object[data.length + 1];
rowData[0] = id;
System.arraycopy(data, 0, rowData, 1, data.length);
dataTableModel.addRow(rowData);
}
private void performSearch() {
if (searchField.getForeground().equals(Color.BLACK)) {
RowFilter<Object, Object> rowFilter = new RowFilter<Object, Object>() {
public boolean include(Entry<?, ?> entry) {
String searchFieldTextText = searchField.getText();
Pattern pattern = null;
try {
pattern = Pattern.compile(searchFieldTextText, Pattern.CASE_INSENSITIVE);
} catch (Exception ignored) {
}
String entryValue = ((String) entry.getValue(1)).toLowerCase();
searchFieldTextText = searchFieldTextText.toLowerCase();
if (pattern != null) {
return searchFieldTextText.isEmpty() || pattern.matcher(entryValue).find() != searchMode.isSelected();
} else {
return searchFieldTextText.isEmpty() || entryValue.contains(searchFieldTextText) != searchMode.isSelected();
}
}
};
sorter.setRowFilter(rowFilter);
}
}
public void setTableListener(MessageTableModel messagePanel) {
// 表格复制功能
dataTable.setTransferHandler(new TransferHandler() {
@Override
public void exportToClipboard(JComponent comp, Clipboard clip, int action) throws IllegalStateException {
if (comp instanceof JTable) {
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);
}
}
});
dataTable.setDefaultEditor(Object.class, null);
// 表格内容双击事件
dataTable.addMouseListener(new MouseAdapter() {
@Override
public void mouseClicked(MouseEvent e) {
if (e.getClickCount() == 2) {
int selectedRow = dataTable.getSelectedRow();
if (selectedRow != -1) {
String rowData = dataTable.getValueAt(selectedRow, 1).toString();
messagePanel.applyMessageFilter(tabName, rowData);
}
}
}
});
}
public String getSelectedDataAtTable(JTable table) {
int[] selectRows = table.getSelectedRows();
StringBuilder selectData = new StringBuilder();
for (int row : selectRows) {
selectData.append(table.getValueAt(row, 1).toString()).append("\n");
}
if (!selectData.isEmpty()) {
selectData.deleteCharAt(selectData.length() - 1);
} else {
return "";
}
return selectData.toString();
}
public JTable getDataTable() {
return this.dataTable;
}
}

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

@@ -0,0 +1,157 @@
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);
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);
}
}

347
src/main/java/hae/component/board/table/Datatable.java Normal file
View File

@@ -0,0 +1,347 @@
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;
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.Comparator;
import java.util.List;
import java.util.Map;
import java.util.regex.Pattern;
public class Datatable extends JPanel {
private final MontoyaApi api;
private final ConfigLoader configLoader;
private final JTable dataTable;
private final DefaultTableModel dataTableModel;
private final JTextField searchField;
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);
this.dataTable = new JTable(dataTableModel);
this.sorter = new TableRowSorter<>(dataTableModel);
this.searchField = new JTextField();
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
public int compare(Integer s1, Integer s2) {
return s1.compareTo(s2);
}
});
dataTable.setRowSorter(sorter);
TableColumn idColumn = dataTable.getColumnModel().getColumn(0);
idColumn.setMaxWidth(50);
for (String item : dataList) {
if (!item.isEmpty()) {
addRowToTable(new Object[]{item});
}
}
// 设置灰色默认文本
String searchText = "Search";
UIEnhancer.setTextFieldPlaceholder(searchField, searchText);
// 监听输入框内容输入、更新、删除
searchField.getDocument().addDocumentListener(new DocumentListener() {
@Override
public void insertUpdate(DocumentEvent e) {
performSearch();
}
@Override
public void removeUpdate(DocumentEvent e) {
performSearch();
}
@Override
public void changedUpdate(DocumentEvent e) {
performSearch();
}
});
// 设置布局
JScrollPane scrollPane = new JScrollPane(dataTable);
scrollPane.setVerticalScrollBarPolicy(ScrollPaneConstants.VERTICAL_SCROLLBAR_ALWAYS);
setLayout(new BorderLayout(0, 5));
JPanel optionsPanel = new JPanel();
optionsPanel.setLayout(new BoxLayout(optionsPanel, BoxLayout.X_AXIS));
// Settings按钮
JPanel settingMenuPanel = new JPanel(new GridLayout(1, 1));
settingMenuPanel.setBorder(BorderFactory.createEmptyBorder(3, 3, 3, 3));
JPopupMenu settingMenu = new JPopupMenu();
settingMenuPanel.add(searchMode);
searchMode.addItemListener(e -> performSearch());
settingMenu.add(settingMenuPanel);
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(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) {
button.addActionListener(e -> {
Point buttonLocation = button.getLocationOnScreen();
Dimension menuSize = menu.getPreferredSize();
int x = buttonLocation.x + (button.getWidth() - menuSize.width) / 2;
int y = buttonLocation.y - menuSize.height;
menu.show(button, x - buttonLocation.x, y - buttonLocation.y);
});
}
private void setProgressBar(boolean status) {
Databoard.setProgressBar(status, progressBar, "AI+ ...");
}
private void addRowToTable(Object[] data) {
int rowCount = dataTableModel.getRowCount();
int id = rowCount > 0 ? (Integer) dataTableModel.getValueAt(rowCount - 1, 0) + 1 : 1;
Object[] rowData = new Object[data.length + 1];
rowData[0] = id;
System.arraycopy(data, 0, rowData, 1, data.length);
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() {
if (searchField.getForeground().equals(Color.BLACK)) {
RowFilter<Object, Object> rowFilter = new RowFilter<Object, Object>() {
public boolean include(Entry<?, ?> entry) {
String searchFieldTextText = searchField.getText();
Pattern pattern = null;
try {
pattern = Pattern.compile(searchFieldTextText, Pattern.CASE_INSENSITIVE);
} catch (Exception ignored) {
}
String entryValue = ((String) entry.getValue(1)).toLowerCase();
searchFieldTextText = searchFieldTextText.toLowerCase();
if (pattern != null) {
return searchFieldTextText.isEmpty() || pattern.matcher(entryValue).find() != searchMode.isSelected();
} else {
return searchFieldTextText.isEmpty() || entryValue.contains(searchFieldTextText) != searchMode.isSelected();
}
}
};
sorter.setRowFilter(rowFilter);
}
}
public void setTableListener(MessageTableModel messagePanel) {
// 表格复制功能
dataTable.setTransferHandler(new TransferHandler() {
@Override
public void exportToClipboard(JComponent comp, Clipboard clip, int action) throws IllegalStateException {
if (comp instanceof JTable) {
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);
}
}
});
dataTable.setDefaultEditor(Object.class, null);
// 表格内容双击事件
dataTable.addMouseListener(new MouseAdapter() {
@Override
public void mouseClicked(MouseEvent e) {
if (e.getClickCount() == 2) {
int selectedRow = dataTable.getSelectedRow();
if (selectedRow != -1) {
String rowData = dataTable.getValueAt(selectedRow, 1).toString();
messagePanel.applyMessageFilter(tabName, rowData);
}
}
}
});
}
private String getTableData(JTable table) {
StringBuilder selectData = new StringBuilder();
int rowCount = table.getRowCount();
for (int i = 0; i < rowCount; i++) {
selectData.append(table.getValueAt(i, 1).toString()).append("\r\n");
}
if (!selectData.isEmpty()) {
selectData.delete(selectData.length() - 2, selectData.length());
} else {
return "";
}
return selectData.toString();
}
public String getSelectedDataAtTable(JTable table) {
int[] selectRows = table.getSelectedRows();
StringBuilder selectData = new StringBuilder();
for (int row : selectRows) {
selectData.append(table.getValueAt(row, 1).toString()).append("\r\n");
}
if (!selectData.isEmpty()) {
selectData.delete(selectData.length() - 2, selectData.length());
} else {
return "";
}
return selectData.toString();
}
public JTable getDataTable() {
return this.dataTable;
}
}

235
src/main/java/hae/component/config/Config.java
View File

@@ -7,6 +7,9 @@ import hae.utils.UIEnhancer;
import javax.swing.*;
import javax.swing.border.EmptyBorder;
import javax.swing.border.TitledBorder;
import javax.swing.event.DocumentEvent;
import javax.swing.event.DocumentListener;
import javax.swing.event.TableModelEvent;
import javax.swing.event.TableModelListener;
import javax.swing.table.DefaultTableModel;
@@ -14,12 +17,11 @@ import java.awt.*;
import java.awt.datatransfer.Clipboard;
import java.awt.datatransfer.DataFlavor;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.awt.event.KeyAdapter;
import java.awt.event.KeyEvent;
import java.util.ArrayList;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Set;
import java.util.*;
public class Config extends JPanel {
private final MontoyaApi api;
@@ -27,22 +29,22 @@ public class Config extends JPanel {
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));
@@ -62,6 +64,153 @@ public class Config extends JPanel {
reloadButton.addActionListener(this::reloadActionPerformed);
updateButton.addActionListener(this::onlineUpdateActionPerformed);
constraints.gridx = 1;
JTabbedPane configTabbedPanel = new JTabbedPane();
String[] settingMode = new String[]{"Exclude suffix", "Block host"};
JPanel settingPanel = createConfigTablePanel(settingMode, "Setting");
JPanel scopePanel = getScopePanel();
JScrollPane scopeScrollPane = new JScrollPane(scopePanel);
scopeScrollPane.setBorder(new TitledBorder("Scope"));
settingPanel.add(scopeScrollPane, 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);
}
private JPanel getScopePanel() {
JPanel scopePanel = new JPanel();
scopePanel.setLayout(new BoxLayout(scopePanel, BoxLayout.X_AXIS));
String[] scopeInit = hae.Config.scopeOptions.split("\\|");
String[] scopeMode = configLoader.getScope().split("\\|");
for (String scope : scopeInit) {
JCheckBox checkBox = new JCheckBox(scope);
scopePanel.add(checkBox);
for (String mode : scopeMode) {
if (scope.equals(mode)) {
checkBox.setSelected(true);
}
}
checkBox.addActionListener(e -> updateScope(checkBox));
}
return scopePanel;
}
private TableModelListener craeteSettingTableModelListener(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("Exclude suffix")) {
if (!values.equals(configLoader.getExcludeSuffix()) && !values.isEmpty()) {
configLoader.setExcludeSuffix(values);
}
}
if (selected.equals("Block host")) {
if (!values.equals(configLoader.getBlockHost()) && !values.isEmpty()) {
configLoader.setBlockHost(values);
}
}
}
};
}
private ActionListener createSettingActionListener(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("Exclude suffix")) {
addDataToTable(configLoader.getExcludeSuffix().replaceAll("\\|", "\r\n"), model);
}
if (selected.equals("Block host")) {
addDataToTable(configLoader.getBlockHost().replaceAll("\\|", "\r\n"), model);
}
}
};
}
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) {
JPanel settingPanel = new JPanel(new BorderLayout());
DefaultTableModel model = new DefaultTableModel();
@@ -80,48 +229,19 @@ public class Config extends JPanel {
JPanel inputPanelB = new JPanel(new BorderLayout());
inputPanelB.setBorder(new EmptyBorder(0, 0, 3, 0));
constraints.gridx = 1;
JButton addButton = new JButton("Add");
JButton removeButton = new JButton("Remove");
JButton pasteButton = new JButton("Paste");
JButton clearButton = new JButton("Clear");
JComboBox<String> setTypeComboBox = new JComboBox<>();
String[] mode = new String[]{"Exclude suffix", "Block host"};
setTypeComboBox.setModel(new DefaultComboBoxModel<>(mode));
setTypeComboBox.addActionListener(e -> {
String selected = (String) setTypeComboBox.getSelectedItem();
model.setRowCount(0);
if (selected.equals("Exclude suffix")) {
addDataToTable(configLoader.getExcludeSuffix().replaceAll("\\|", "\r\n"), model);
}
setTypeComboBox.addActionListener(type.equals("AI+") ? createAIActionListener(setTypeComboBox, model) : createSettingActionListener(setTypeComboBox, model));
if (selected.equals("Block host")) {
addDataToTable(configLoader.getBlockHost().replaceAll("\\|", "\r\n"), model);
}
});
setTypeComboBox.setSelectedItem("Exclude suffix");
setTypeComboBox.setSelectedItem(mode[0]);
model.addTableModelListener(new TableModelListener() {
@Override
public void tableChanged(TableModelEvent e) {
String selected = (String) setTypeComboBox.getSelectedItem();
String values = getFirstColumnDataAsString(model);
if (selected.equals("Exclude suffix")) {
if (!values.equals(configLoader.getExcludeSuffix()) && !values.isEmpty()) {
configLoader.setExcludeSuffix(values);
}
}
if (selected.equals("Block host")) {
if (!values.equals(configLoader.getBlockHost()) && !values.isEmpty()) {
configLoader.setBlockHost(values);
}
}
}
});
model.addTableModelListener(type.equals("AI+") ? craeteAITableModelListener(setTypeComboBox, model) : craeteSettingTableModelListener(setTypeComboBox, model));
constraints.insets = new Insets(0, 0, 3, 0);
constraints.gridy = 0;
@@ -146,13 +266,13 @@ public class Config extends JPanel {
settingPanel.add(inputPanel, BorderLayout.CENTER);
addButton.addActionListener(e -> addActionPerformedAction(e, model));
addButton.addActionListener(e -> addActionPerformed(e, model));
addTextField.addKeyListener(new KeyAdapter() {
@Override
public void keyPressed(KeyEvent e) {
if (e.getKeyCode() == KeyEvent.VK_ENTER) {
addActionPerformedAction(null, model);
addActionPerformed(null, model);
}
}
});
@@ -178,16 +298,14 @@ public class Config extends JPanel {
clearButton.addActionListener(e -> model.setRowCount(0));
JPanel settingMainPanel = new JPanel(new BorderLayout());
JLabel settingLabel = new JLabel("Setting:");
JPanel settingLabelPanel = new JPanel(new BorderLayout());
settingLabelPanel.add(settingLabel, BorderLayout.WEST);
settingMainPanel.setBorder(new EmptyBorder(0, 15, 10, 15));
settingMainPanel.add(settingLabelPanel, BorderLayout.NORTH);
settingMainPanel.add(settingPanel, BorderLayout.CENTER);
add(ruleInfoPanel, BorderLayout.NORTH);
add(settingMainPanel, BorderLayout.CENTER);
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"));
settingMainPanel.add(settingScroller, BorderLayout.CENTER);
return settingMainPanel;
}
@@ -239,7 +357,22 @@ public class Config extends JPanel {
}
}
private void addActionPerformedAction(ActionEvent e, DefaultTableModel model) {
public void updateScope(JCheckBox checkBox) {
String boxText = checkBox.getText();
boolean selected = checkBox.isSelected();
Set<String> HaEScope = new HashSet<>(Arrays.asList(configLoader.getScope().split("\\|")));
if (selected) {
HaEScope.add(boxText);
} else {
HaEScope.remove(boxText);
}
configLoader.setScope(String.join("|", HaEScope));
}
private void addActionPerformed(ActionEvent e, DefaultTableModel model) {
String addTextFieldText = addTextField.getText();
if (!addTextFieldText.equals(defaultText)) {
addDataToTable(addTextFieldText, model);