package burp.datacollector.dao; import com.opencsv.CSVWriter; import java.io.File; import java.io.FileOutputStream; import java.io.FileWriter; import java.io.IOException; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; import java.util.HashSet; public class HostFileMapDao extends BaseDao { public final static String FILE_IMPORT_FILE = "/file_import.csv"; public final static String FILE_FILE = "/file.txt"; public void insertIgnoreHostFile(String host, HashSet<String> fileNameSet) throws SQLException { StringBuilder sqlStringBuilder = new StringBuilder("INSERT IGNORE INTO host_file_map(host, filename) VALUES"); int n = fileNameSet.size(); for (int i = 0; i < n - 1; i++) sqlStringBuilder.append("(?,?), "); sqlStringBuilder.append("(?,?)"); String sql = sqlStringBuilder.toString(); PreparedStatement preparedStatement = getPreparedStatement(sql); int length = 1; int hostIndex = 1; int index = 2; for (String fileName : fileNameSet) { preparedStatement.setString(hostIndex, host); preparedStatement.setString(index, fileName); length += 2; hostIndex = length; index = hostIndex + 1; } preparedStatement.executeUpdate(); preparedStatement.close(); } public void exportFile(String dirName, int fileCount) throws SQLException, IOException { String sql = "SELECT stat.filename, sum(fileCount) AS allCount\n" + "FROM ((SELECT hfm.filename, count(*) AS fileCount FROM host_file_map hfm GROUP BY hfm.filename)\n" + " UNION ALL\n" + " (SELECT filename, count AS fileCount FROM file)) stat\n" + "GROUP BY stat.filename\n" + "HAVING allCount >= ?\n" + "ORDER BY allCount DESC"; PreparedStatement preparedStatement = getPreparedStatement(sql); preparedStatement.setInt(1, fileCount); ResultSet resultSet = preparedStatement.executeQuery(); File file = new File(dirName + FILE_FILE); File fileImportFile = new File(dirName + FILE_IMPORT_FILE); FileOutputStream fileOutputStream = new FileOutputStream(file); FileWriter fileWriter = new FileWriter(fileImportFile); CSVWriter csvWriter = new CSVWriter(fileWriter); String[] fileHead = new String[]{"file", "count"}; csvWriter.writeNext(fileHead); while (resultSet.next()) { String fileName = resultSet.getString(1); String row = fileName + "\n"; int count = resultSet.getInt(2); fileOutputStream.write(row.getBytes()); csvWriter.writeNext(new String[]{fileName, String.valueOf(count)}, true); } fileOutputStream.close(); csvWriter.close(); } }