400 lines
15 KiB
Java
400 lines
15 KiB
Java
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.ArrayList;
|
|
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 JTextField secondSearchField;
|
|
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(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
|
|
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});
|
|
}
|
|
}
|
|
|
|
UIEnhancer.setTextFieldPlaceholder(searchField, "Search");
|
|
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();
|
|
}
|
|
|
|
});
|
|
|
|
UIEnhancer.setTextFieldPlaceholder(secondSearchField, "Second search");
|
|
secondSearchField.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(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) {
|
|
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() {
|
|
RowFilter<Object, Object> firstRowFilter = applyFirstSearchFilter();
|
|
RowFilter<Object, Object> secondRowFilter = applySecondFilter();
|
|
if (searchField.getForeground().equals(Color.BLACK)) {
|
|
sorter.setRowFilter(firstRowFilter);
|
|
if (secondSearchField.getForeground().equals(Color.BLACK)) {
|
|
List<RowFilter<Object, Object>> filters = new ArrayList<>();
|
|
filters.add(firstRowFilter);
|
|
filters.add(secondRowFilter);
|
|
sorter.setRowFilter(RowFilter.andFilter(filters));
|
|
}
|
|
}
|
|
}
|
|
|
|
private RowFilter<Object, Object> applyFirstSearchFilter() {
|
|
return 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();
|
|
}
|
|
}
|
|
};
|
|
}
|
|
|
|
private RowFilter<Object, Object> applySecondFilter() {
|
|
return new RowFilter<Object, Object>() {
|
|
public boolean include(Entry<?, ?> entry) {
|
|
String searchFieldTextText = secondSearchField.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();
|
|
} else {
|
|
return searchFieldTextText.isEmpty() || entryValue.contains(searchFieldTextText);
|
|
}
|
|
}
|
|
};
|
|
}
|
|
|
|
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;
|
|
}
|
|
}
|
|
|