diff --git a/src/main/java/caa/CaA.java b/src/main/java/caa/CaA.java index e0c1936..05e84a3 100644 --- a/src/main/java/caa/CaA.java +++ b/src/main/java/caa/CaA.java @@ -2,6 +2,7 @@ import burp.api.montoya.BurpExtension; import burp.api.montoya.MontoyaApi; +import burp.api.montoya.extension.ExtensionUnloadingHandler; import burp.api.montoya.logging.Logging; import caa.component.CaAMain; import caa.component.member.taskboard.MessageTableModel; @@ -12,12 +13,13 @@ import java.io.File; import java.sql.Connection; +import java.sql.SQLException; public class CaA implements BurpExtension { @Override public void initialize(MontoyaApi api) { // 设置扩展名称 - String version = "Beta 0.8"; + String version = "Beta 0.9"; Database db = null; MessageTableModel messageTableModel = new MessageTableModel(api); api.extension().setName(String.format("CaA (%s) - Collector and Analyzer", version)); @@ -59,5 +61,19 @@ public void initialize(MontoyaApi api) { api.logging().logToOutput("[Error] Failed to connect to the CaA database!"); api.extension().unload(); } + + api.extension().registerUnloadingHandler(new ExtensionUnloadingHandler() { + @Override + public void extensionUnloaded() { + try { + if (con != null) { + con.close(); + } + } catch (Exception ignored) { + } + + } + }); + } } diff --git a/src/main/java/caa/component/member/FuzzerDialog.java b/src/main/java/caa/component/member/FuzzerDialog.java index 05c0085..3abc7e2 100644 --- a/src/main/java/caa/component/member/FuzzerDialog.java +++ b/src/main/java/caa/component/member/FuzzerDialog.java @@ -17,6 +17,8 @@ import java.awt.datatransfer.DataFlavor; import java.awt.event.ComponentAdapter; import java.awt.event.ComponentEvent; +import java.awt.event.MouseAdapter; +import java.awt.event.MouseEvent; import java.util.*; import java.util.List; import java.util.stream.IntStream; @@ -37,6 +39,15 @@ public class FuzzerDialog extends JDialog { private final String payload; private final boolean secure; private final Dimension dialogDimension = new Dimension(888, 688); + private final JPopupMenu popupMenu; + + private final MouseAdapter mouseAdapter = new MouseAdapter() { + public void mouseReleased(MouseEvent e) { + if (SwingUtilities.isRightMouseButton(e)) { + popupMenu.show(e.getComponent(), e.getX(), e.getY()); + } + } + }; public FuzzerDialog(MontoyaApi api, Database db, MessageTableModel messageTableModel, String httpRequest, boolean secure, String tabName, String payload) { this.api = api; @@ -48,6 +59,8 @@ public FuzzerDialog(MontoyaApi api, Database db, MessageTableModel messageTableM this.secure = secure; contentPanel = new JPanel(new BorderLayout()); + popupMenu = new JPopupMenu(); + initComponents(); setTitle("Fuzzer - Task Configuration"); @@ -145,13 +158,32 @@ private void updateHost() { JComboBox fuzzModeComboBox = new JComboBox<>(); String[] mode = new String[] {"Param", "File", "Path", "Value"}; fuzzModeComboBox.setModel(new DefaultComboBoxModel<>(mode)); + + // Value 修改Name快捷方式 + + JMenuItem renameMenuItem = new JMenuItem("Rename"); + popupMenu.add(renameMenuItem); + renameMenuItem.addActionListener(e -> { + int[] selectedRows = table.getSelectedRows(); + String newName = JOptionPane.showInputDialog("Please enter the new name."); + if (!newName.isBlank()) { + for (int row : selectedRows) { + int columnIndex = 0; + table.getModel().setValueAt(newName, row, columnIndex); + } + } + }); + fuzzModeComboBox.addActionListener(e -> { String selected = (String) fuzzModeComboBox.getSelectedItem(); if (selected.equals("Value")) { model.addColumn("Value"); model.setColumnCount(2); + + table.addMouseListener(mouseAdapter); } else { model.setColumnCount(1); + table.removeMouseListener(mouseAdapter); } addDataToTable(payload, model); }); diff --git a/src/main/java/caa/instances/Database.java b/src/main/java/caa/instances/Database.java index 7d90a16..192c4a8 100644 --- a/src/main/java/caa/instances/Database.java +++ b/src/main/java/caa/instances/Database.java @@ -19,17 +19,25 @@ public Database(MontoyaApi api, String dbFileName) { try { Class.forName("org.sqlite.JDBC"); this.connection = DriverManager.getConnection(String.format("jdbc:sqlite:%s", dbFileName)); - // 初始化表 - createTables(); - } catch (Exception e) { - api.logging().logToError(e); - } - } + Statement statement = this.connection.createStatement(); - private void executeSql(String sql) { - try { - Statement stmt = connection.createStatement(); - stmt.execute(sql); + // 开启 WAL 模式 + statement.executeUpdate("PRAGMA journal_mode=WAL;"); + + // 修改分页大小 + statement.executeUpdate("PRAGMA page_size = 8192;"); + + // 开启自动检查点模式 + statement.executeUpdate("PRAGMA wal_autocheckpoint = 2000;"); + + // 关闭同步模式 + statement.executeUpdate("PRAGMA synchronous=OFF;"); + + // 启用缓存 + statement.executeUpdate("PRAGMA cache_size = 10000;"); + + // 初始化数据表 + createTables(); } catch (Exception e) { api.logging().logToError(e); } @@ -37,57 +45,59 @@ private void executeSql(String sql) { public Object selectData(String host, String tableName, String limitSize) { try { - String sql = "select name%s from `" + tableName + "` %s order by count desc"; - if (!limitSize.isBlank()) { - sql += " limit " + limitSize; - } - if (tableName.contains("All")) { - sql = String.format(sql, "", ""); - } else if (tableName.equals("Value")){ - sql = String.format(sql, ",value", "where host = ?"); - } else { - sql = String.format(sql, "", "where host = ?"); - } + if (!connection.isClosed()) { + String sql = "select name%s from `" + tableName + "` %s order by count desc"; + if (!limitSize.isBlank()) { + sql += " limit " + limitSize; + } + if (tableName.contains("All")) { + sql = String.format(sql, "", ""); + } else if (tableName.equals("Value")){ + sql = String.format(sql, ",value", "where host = ?"); + } else { + sql = String.format(sql, "", "where host = ?"); + } - PreparedStatement ps = connection.prepareStatement(sql); - prepareStatement(host, ps); - ResultSet rs = ps.executeQuery(); + PreparedStatement ps = connection.prepareStatement(sql); + prepareStatement(host, ps); + ResultSet rs = ps.executeQuery(); - // 判断结果集是否为空 - if (!rs.isBeforeFirst()) { - return null; - } else { - if (tableName.equals("Value")) { - SetMultimap multimap = LinkedHashMultimap.create(); - while (rs.next()){ - String key = rs.getString(1); - String value = rs.getString(2); - multimap.put(key, value); - } - if (multimap.size() <= 0) { - return null; - } - return multimap; + // 判断结果集是否为空 + if (!rs.isBeforeFirst()) { + return null; } else { - Set resultList = new HashSet<>(); - while (rs.next()){ - String columnValue = rs.getString(1); - resultList.add(columnValue); - } - if (resultList.isEmpty()) { - return null; + if (tableName.equals("Value")) { + SetMultimap multimap = LinkedHashMultimap.create(); + while (rs.next()){ + String key = rs.getString(1); + String value = rs.getString(2); + multimap.put(key, value); + } + if (multimap.size() <= 0) { + return null; + } + return multimap; + } else { + Set resultList = new LinkedHashSet<>(); + while (rs.next()){ + String columnValue = rs.getString(1); + resultList.add(columnValue); + } + if (resultList.isEmpty()) { + return null; + } + return resultList; } - return resultList; } } } catch (Exception e) { api.logging().logToError(e); - return null; } + return null; } - public void createTables() { + private void createTables() { String sqlTemplate = """ CREATE TABLE IF NOT EXISTS `%s` ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -109,8 +119,12 @@ name TEXT(300) DEFAULT NULL, } else { uniqueField = ", UNIQUE(name)"; } - String sql = String.format(sqlTemplate, name, fields, uniqueField); - executeSql(sql); + try { + Statement stmt = connection.createStatement(); + stmt.execute(String.format(sqlTemplate, name, fields, uniqueField)); + } catch (Exception e) { + api.logging().logToError(e); + } } }