Version: 3.3 Update
This commit is contained in:
gh0stkey
347
src/main/java/hae/component/board/table/Datatable.java
Normal file
347
src/main/java/hae/component/board/table/Datatable.java
Normal 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;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user