From 6d4abae89810dda437319800b5726d86e7010df9 Mon Sep 17 00:00:00 2001 From: gh0stkey <24655118+gh0stkey@users.noreply.github.com> Date: Thu, 30 May 2024 14:37:01 +0800 Subject: [PATCH] Version: 3.2.1 Update --- src/main/java/hae/Config.java | 2 + src/main/java/hae/HaE.java | 4 +- .../java/hae/component/board/Databoard.java | 410 ++++++++++++------ .../component/board/message/MessageEntry.java | 14 +- .../board/message/MessageTableModel.java | 196 ++++++--- .../java/hae/component/config/Config.java | 2 +- src/main/java/hae/component/rule/Rule.java | 6 +- src/main/java/hae/component/rule/Rules.java | 2 +- .../hae/instances/editor/RequestEditor.java | 4 +- .../instances/http/HttpMessageHandler.java | 7 +- .../instances/http/utils/RegularMatcher.java | 2 +- .../java/hae/utils/project/FileProcessor.java | 47 ++ .../hae/utils/project/ProjectProcessor.java | 173 ++++++-- .../utils/project/model/HaeFileContent.java | 23 +- .../hae/utils/string/StringProcessor.java | 42 ++ 15 files changed, 695 insertions(+), 239 deletions(-) create mode 100644 src/main/java/hae/utils/project/FileProcessor.java diff --git a/src/main/java/hae/Config.java b/src/main/java/hae/Config.java index b255169..1ce6d7e 100644 --- a/src/main/java/hae/Config.java +++ b/src/main/java/hae/Config.java @@ -54,4 +54,6 @@ public class Config { public static Map globalRules = new HashMap<>(); public static ConcurrentHashMap>> globalDataMap = new ConcurrentHashMap<>(); + + public static ConcurrentHashMap> globalHostHashMap = new ConcurrentHashMap<>(); } diff --git a/src/main/java/hae/HaE.java b/src/main/java/hae/HaE.java index 937c7b1..c0f408b 100644 --- a/src/main/java/hae/HaE.java +++ b/src/main/java/hae/HaE.java @@ -16,13 +16,13 @@ public class HaE implements BurpExtension { @Override public void initialize(MontoyaApi api) { // 设置扩展名称 - String version = "3.2"; + String version = "3.2.1"; api.extension().setName(String.format("HaE (%s) - Highlighter and Extractor", version)); // 加载扩展后输出的项目信息 Logging logging = api.logging(); logging.logToOutput("[ HACK THE WORLD - TO DO IT ]"); - logging.logToOutput("[#] Author: EvilChen && 0chencc"); + logging.logToOutput("[#] Author: EvilChen && 0chencc && vaycore"); logging.logToOutput("[#] Github: https://github.com/gh0stkey/HaE"); // 配置文件加载 diff --git a/src/main/java/hae/component/board/Databoard.java b/src/main/java/hae/component/board/Databoard.java index 2bdeedf..f491a14 100644 --- a/src/main/java/hae/component/board/Databoard.java +++ b/src/main/java/hae/component/board/Databoard.java @@ -1,10 +1,6 @@ package hae.component.board; import burp.api.montoya.MontoyaApi; -import burp.api.montoya.http.HttpService; -import burp.api.montoya.http.message.HttpRequestResponse; -import burp.api.montoya.http.message.requests.HttpRequest; -import burp.api.montoya.http.message.responses.HttpResponse; import hae.Config; import hae.component.board.message.MessageEntry; import hae.component.board.message.MessageTableModel; @@ -27,7 +23,8 @@ import java.awt.event.*; import java.io.File; import java.util.List; import java.util.*; -import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.*; +import java.util.function.Function; import java.util.stream.Collectors; public class Databoard extends JPanel { @@ -35,15 +32,19 @@ public class Databoard extends JPanel { private final ConfigLoader configLoader; private final ProjectProcessor projectProcessor; private final MessageTableModel messageTableModel; + private JTextField hostTextField; private JTabbedPane dataTabbedPane; private JSplitPane splitPane; private MessageTable messageTable; + private JProgressBar progressBar; private static Boolean isMatchHost = false; private final DefaultComboBoxModel comboBoxModel = new DefaultComboBoxModel(); private final JComboBox hostComboBox = new JComboBox(comboBoxModel); + private SwingWorker currentWorker; + public Databoard(MontoyaApi api, ConfigLoader configLoader, MessageTableModel messageTableModel) { this.api = api; this.configLoader = configLoader; @@ -56,9 +57,9 @@ 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, 0}; + ((GridBagLayout) getLayout()).rowHeights = new int[]{0, 65, 20, 25, 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, 1.0E-4}; + ((GridBagLayout) getLayout()).rowWeights = new double[]{0.0, 1.0, 0.0, 0.0, 1.0E-4}; JLabel hostLabel = new JLabel("Host:"); @@ -78,6 +79,10 @@ public class Databoard extends JPanel { splitPane = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT); dataTabbedPane = new JTabbedPane(JTabbedPane.TOP); + dataTabbedPane.setPreferredSize(new Dimension(500, 0)); + dataTabbedPane.setTabLayoutPolicy(JTabbedPane.SCROLL_TAB_LAYOUT); + splitPane.setLeftComponent(dataTabbedPane); + actionButton.addActionListener(e -> { int x = 0; int y = actionButton.getHeight(); @@ -88,6 +93,8 @@ public class Databoard extends JPanel { exportButton.addActionListener(this::exportActionPerformed); importButton.addActionListener(this::importActionPerformed); + progressBar = new JProgressBar(); + splitPane.addComponentListener(new ComponentAdapter() { @Override public void componentResized(ComponentEvent e) { @@ -96,6 +103,7 @@ public class Databoard extends JPanel { }); splitPane.setVisible(false); + progressBar.setVisible(false); add(hostLabel, new GridBagConstraints(1, 0, 1, 1, 0.0, 0.0, GridBagConstraints.CENTER, GridBagConstraints.BOTH, new Insets(8, 0, 5, 5), 0, 0)); @@ -103,9 +111,13 @@ public class Databoard extends JPanel { new Insets(8, 0, 5, 5), 0, 0)); add(actionButton, new GridBagConstraints(3, 0, 1, 1, 0.0, 0.0, GridBagConstraints.CENTER, GridBagConstraints.BOTH, new Insets(8, 0, 5, 5), 0, 0)); - add(splitPane, new GridBagConstraints(1, 1, 3, 3, 0.0, 0.0, + + add(splitPane, new GridBagConstraints(1, 1, 3, 1, 0.0, 1.0, GridBagConstraints.CENTER, GridBagConstraints.BOTH, - new Insets(8, 0, 5, 5), 0, 0)); + new Insets(0, 5, 0, 5), 0, 0)); + add(progressBar, new GridBagConstraints(1, 3, 3, 1, 1.0, 0.0, + GridBagConstraints.CENTER, GridBagConstraints.HORIZONTAL, + new Insets(0, 5, 0, 5), 0, 0)); hostComboBox.setMaximumRowCount(5); add(hostComboBox, new GridBagConstraints(2, 0, 1, 1, 0.0, 0.0, GridBagConstraints.CENTER, GridBagConstraints.BOTH, new Insets(8, 0, 5, 5), 0, 0)); @@ -125,6 +137,19 @@ public class Databoard extends JPanel { columnModel.getColumn(5).setPreferredWidth((int) (totalWidth * 0.1)); } + private void setProgressBar(boolean status) { + progressBar.setIndeterminate(status); + if (!status) { + progressBar.setMaximum(100); + progressBar.setString("OK"); + progressBar.setStringPainted(true); + progressBar.setValue(progressBar.getMaximum()); + } else { + progressBar.setString("Loading..."); + progressBar.setStringPainted(true); + } + } + private void setAutoMatch() { hostComboBox.setSelectedItem(null); hostComboBox.addActionListener(this::handleComboBoxAction); @@ -157,9 +182,45 @@ public class Databoard extends JPanel { private void handleComboBoxAction(ActionEvent e) { if (!isMatchHost && hostComboBox.getSelectedItem() != null) { + progressBar.setVisible(true); + setProgressBar(true); String selectedHost = hostComboBox.getSelectedItem().toString(); hostTextField.setText(selectedHost); - populateTabbedPaneByHost(selectedHost); + + if (currentWorker != null && !currentWorker.isDone()) { + currentWorker.cancel(true); + } + + currentWorker = new SwingWorker() { + @Override + protected Boolean doInBackground() { + return populateTabbedPaneByHost(selectedHost); + } + + @Override + protected void done() { + if (!isCancelled()) { + try { + boolean status = get(); + if (status) { + JSplitPane messageSplitPane = messageTableModel.getSplitPane(); + splitPane.setRightComponent(messageSplitPane); + messageTable = messageTableModel.getMessageTable(); + resizePanel(); + + splitPane.setVisible(true); + hostTextField.setText(selectedHost); + + hostComboBox.setPopupVisible(false); + applyHostFilter(selectedHost); + } + } catch (Exception ignored) { + } + } + } + }; + + currentWorker.execute(); } } @@ -178,7 +239,6 @@ public class Databoard extends JPanel { if (keyCode == KeyEvent.VK_ENTER) { isMatchHost = false; handleComboBoxAction(null); - hostComboBox.setPopupVisible(false); } if (keyCode == KeyEvent.VK_ESCAPE) { @@ -188,10 +248,52 @@ public class Databoard extends JPanel { isMatchHost = false; } + private boolean populateTabbedPaneByHost(String selectedHost) { + ConcurrentHashMap>> dataMap = Config.globalDataMap; + Map> selectedDataMap; + + if (selectedHost.contains("*")) { + selectedDataMap = new HashMap<>(); + dataMap.keySet().parallelStream().forEach(key -> { + if ((StringProcessor.matchesHostPattern(key, selectedHost) || selectedHost.equals("*")) && !key.contains("*")) { + Map> ruleMap = dataMap.get(key); + for (String ruleKey : ruleMap.keySet()) { + List dataList = ruleMap.get(ruleKey); + if (selectedDataMap.containsKey(ruleKey)) { + List mergedList = new ArrayList<>(selectedDataMap.get(ruleKey)); + mergedList.addAll(dataList); + HashSet uniqueSet = new HashSet<>(mergedList); + selectedDataMap.put(ruleKey, new ArrayList<>(uniqueSet)); + } else { + selectedDataMap.put(ruleKey, dataList); + } + } + } + }); + } else { + selectedDataMap = dataMap.get(selectedHost); + } + + if (!selectedDataMap.isEmpty()) { + dataTabbedPane.removeAll(); + + for (Map.Entry> 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; + } + private void filterComboBoxList() { isMatchHost = true; comboBoxModel.removeAllElements(); String input = hostTextField.getText().toLowerCase(); + if (!input.isEmpty()) { for (String host : getHostByList()) { String lowerCaseHost = host.toLowerCase(); @@ -210,83 +312,40 @@ public class Databoard extends JPanel { isMatchHost = false; } - private void populateTabbedPaneByHost(String selectedHost) { - if (!Objects.equals(selectedHost, "")) { - ConcurrentHashMap>> dataMap = Config.globalDataMap; - Map> selectedDataMap; - - dataTabbedPane.removeAll(); - dataTabbedPane.setPreferredSize(new Dimension(500, 0)); - dataTabbedPane.setTabLayoutPolicy(JTabbedPane.SCROLL_TAB_LAYOUT); - splitPane.setLeftComponent(dataTabbedPane); - - if (selectedHost.contains("*")) { - // 通配符数据 - selectedDataMap = new HashMap<>(); - for (String key : dataMap.keySet()) { - if ((StringProcessor.matchesHostPattern(key, selectedHost) || selectedHost.equals("*")) && !key.contains("*")) { - Map> ruleMap = dataMap.get(key); - for (String ruleKey : ruleMap.keySet()) { - List dataList = ruleMap.get(ruleKey); - if (selectedDataMap.containsKey(ruleKey)) { - List mergedList = new ArrayList<>(selectedDataMap.get(ruleKey)); - mergedList.addAll(dataList); - HashSet uniqueSet = new HashSet<>(mergedList); - selectedDataMap.put(ruleKey, new ArrayList<>(uniqueSet)); - } else { - selectedDataMap.put(ruleKey, dataList); - } - } - } - } - } else { - selectedDataMap = dataMap.get(selectedHost); - } - - for (Map.Entry> 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); - } - - // 展示请求消息表单 - JSplitPane messageSplitPane = messageTableModel.getSplitPane(); - this.splitPane.setRightComponent(messageSplitPane); - messageTable = messageTableModel.getMessageTable(); - - resizePanel(); - splitPane.setVisible(true); - - applyHostFilter(selectedHost); - hostTextField.setText(selectedHost); - } - } - private void applyHostFilter(String filterText) { TableRowSorter sorter = (TableRowSorter) messageTable.getRowSorter(); - String cleanedText = StringProcessor.replaceFirstOccurrence(filterText, "*.", ""); - RowFilter rowFilter = new RowFilter() { - public boolean include(Entry entry) { - if (cleanedText.equals("*")) { - return true; - } else { - String host = StringProcessor.getHostByUrl((String) entry.getValue(1)); + new SwingWorker() { + @Override + protected Void doInBackground() throws Exception { + RowFilter rowFilter = new RowFilter() { + public boolean include(Entry entry) { + if (cleanedText.equals("*")) { + return true; + } else { + String host = StringProcessor.getHostByUrl((String) entry.getValue(1)); + return StringProcessor.matchesHostPattern(host, filterText); + } + } + }; - return StringProcessor.matchesHostPattern(host, filterText); - } + sorter.setRowFilter(rowFilter); + messageTableModel.applyHostFilter(filterText); + + return null; } - }; - sorter.setRowFilter(rowFilter); + @Override + protected void done() { + setProgressBar(false); + } + }.execute(); - messageTableModel.applyHostFilter(filterText); } private List getHostByList() { - if (!(Config.globalDataMap.keySet().size() == 1 && Config.globalDataMap.keySet().stream().anyMatch(key -> key.contains("*")))) { + if (!Config.globalDataMap.keySet().isEmpty()) { return new ArrayList<>(Config.globalDataMap.keySet()); } return new ArrayList<>(); @@ -305,13 +364,25 @@ public class Databoard extends JPanel { return; } - ConcurrentHashMap>> dataMap = Config.globalDataMap; - List taskStatusList = exportData(selectedHost, exportDir, dataMap); + new SwingWorker, Void>() { + @Override + protected List doInBackground() { + ConcurrentHashMap>> dataMap = Config.globalDataMap; + return exportData(selectedHost, exportDir, dataMap); + } - if (!taskStatusList.isEmpty()) { - String exportStatusMessage = String.format("Exported File List Status:\n%s", String.join("\n", taskStatusList)); - JOptionPane.showConfirmDialog(null, exportStatusMessage, "Info", JOptionPane.YES_OPTION); - } + @Override + protected void done() { + try { + List 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); + } + } catch (Exception ignored) { + } + } + }.execute(); } private List exportData(String selectedHost, String exportDir, Map>> dataMap) { @@ -332,28 +403,63 @@ public class Databoard extends JPanel { } List messageEntryList = messageTableModel.getLogs(); - Map> httpMap = messageEntryList.stream() - .filter(messageEntry -> !StringProcessor.getHostByUrl(messageEntry.getUrl()).isEmpty()) - .filter(messageEntry -> StringProcessor.getHostByUrl(messageEntry.getUrl()).equals(key)) + + Map entryUUIDMap = messageEntryList.stream() .collect(Collectors.toMap( - MessageEntry::getUrl, - this::createHttpItemMap, - (existing, replacement) -> existing + messageEntry -> messageEntry, + messageEntry -> StringProcessor.getRandomUUID(), + (existing, replacement) -> existing // 在冲突时保留现有的映射 )); + Map> httpMap = processEntries( + messageEntryList, + key, + entryUUIDMap, + this::createHttpItemMap + ); + + Map> urlMap = processEntries( + messageEntryList, + key, + entryUUIDMap, + this::creteUrlItemMap + ); + String hostName = key.replace(":", "_"); - String filename = String.format("%s/%s.hae", exportDir, hostName); - boolean createdStatus = projectProcessor.createHaeFile(filename, key, ruleMap, httpMap); + String filename = String.format("%s/%s-%s.hae", exportDir, StringProcessor.getCurrentTime(), hostName); + boolean createdStatus = projectProcessor.createHaeFile(filename, key, ruleMap, urlMap, httpMap); return String.format("Filename: %s, Status: %s", filename, createdStatus); } - private Map createHttpItemMap(MessageEntry entry) { - Map httpItemMap = new HashMap<>(); - httpItemMap.put("comment", entry.getComment()); - httpItemMap.put("color", entry.getColor()); - httpItemMap.put("request", entry.getRequestResponse().request().toString()); - httpItemMap.put("response", entry.getRequestResponse().response().toString()); + + private Map> processEntries(List messageEntryList, String key, Map entryUUIDMap, Function> mapFunction) { + return messageEntryList.stream() + .filter(messageEntry -> !StringProcessor.getHostByUrl(messageEntry.getUrl()).isEmpty()) + .filter(messageEntry -> StringProcessor.getHostByUrl(messageEntry.getUrl()).equals(key)) + .collect(Collectors.toMap( + entryUUIDMap::get, + mapFunction, + (existing, replacement) -> existing + )); + } + + private Map creteUrlItemMap(MessageEntry entry) { + Map urlItemMap = new LinkedHashMap<>(); + urlItemMap.put("url", entry.getUrl()); + urlItemMap.put("method", entry.getMethod()); + urlItemMap.put("status", entry.getStatus()); + urlItemMap.put("length", entry.getLength()); + urlItemMap.put("comment", entry.getComment()); + urlItemMap.put("color", entry.getColor()); + urlItemMap.put("size", String.valueOf(entry.getRequestResponse().request().toByteArray().length())); + return urlItemMap; + } + + private Map createHttpItemMap(MessageEntry entry) { + Map httpItemMap = new LinkedHashMap<>(); + httpItemMap.put("request", entry.getRequestResponse().request().toByteArray().getBytes()); + httpItemMap.put("response", entry.getRequestResponse().response().toByteArray().getBytes()); return httpItemMap; } @@ -363,43 +469,74 @@ public class Databoard extends JPanel { return; } - List filesWithExtension = findFilesWithExtension(new File(exportDir), ".hae"); - List taskStatusList = filesWithExtension.stream() - .map(this::importData) - .collect(Collectors.toList()); + new SwingWorker, Void>() { + @Override + protected List doInBackground() { + List filesWithExtension = findFilesWithExtension(new File(exportDir), ".hae"); + return filesWithExtension.stream() + .map(Databoard.this::importData) + .collect(Collectors.toList()); + } - if (!taskStatusList.isEmpty()) { - String importStatusMessage = "Imported File List Status:\n" + String.join("\n", taskStatusList); - JOptionPane.showConfirmDialog(null, importStatusMessage, "Info", JOptionPane.YES_OPTION); - } + @Override + protected void done() { + try { + List 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); + } + } catch (Exception ignored) { + } + } + }.execute(); } private String importData(String filename) { HaeFileContent haeFileContent = projectProcessor.readHaeFile(filename); boolean readStatus = haeFileContent != null; - if (readStatus) { - String host = haeFileContent.getHost(); - haeFileContent.getDataMap().forEach((key, value) -> RegularMatcher.putDataToGlobalMap(host, key, value)); + ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2); + List> futures = new ArrayList<>(); - haeFileContent.getHttpMap().forEach((key, httpItemMap) -> { - String comment = httpItemMap.get("comment"); - String color = httpItemMap.get("color"); - HttpRequestResponse httpRequestResponse = createHttpRequestResponse(key, httpItemMap); - messageTableModel.add(httpRequestResponse, comment, color); - }); + if (readStatus) { + try { + String host = haeFileContent.getHost(); + haeFileContent.getDataMap().forEach((key, value) -> RegularMatcher.putDataToGlobalMap(host, key, value)); + + haeFileContent.getUrlMap().forEach((key, urlItemMap) -> { + Future future = executor.submit(() -> { + String url = urlItemMap.get("url"); + String comment = urlItemMap.get("comment"); + String color = urlItemMap.get("color"); + String length = urlItemMap.get("length"); + String method = urlItemMap.get("method"); + String status = urlItemMap.get("status"); + String path = haeFileContent.getHttpPath(); + + messageTableModel.add(null, url, method, status, length, comment, color, key, path); + }); + + futures.add(future); + }); + + for (Future future : futures) { + try { + future.get(); + } catch (InterruptedException | ExecutionException e) { + Thread.currentThread().interrupt(); + } + } + } catch (Exception e) { + api.logging().logToError("importData: " + e.getMessage()); + } finally { + executor.shutdown(); + } } return String.format("Filename: %s, Status: %s", filename, readStatus); } - private HttpRequestResponse createHttpRequestResponse(String key, Map httpItemMap) { - HttpService httpService = HttpService.httpService(key); - HttpRequest httpRequest = HttpRequest.httpRequest(httpService, httpItemMap.get("request")); - HttpResponse httpResponse = HttpResponse.httpResponse(httpItemMap.get("response")); - return HttpRequestResponse.httpRequestResponse(httpRequest, httpResponse); - } - private List findFilesWithExtension(File directory, String extension) { List filePaths = new ArrayList<>(); if (directory.isDirectory()) { @@ -413,8 +550,9 @@ public class Databoard extends JPanel { } } } + } else { + filePaths.add(directory.getAbsolutePath()); } - filePaths.add(directory.getAbsolutePath()); return filePaths; } @@ -438,19 +576,35 @@ public class Databoard extends JPanel { } private void clearActionPerformed(ActionEvent e) { - int retCode = JOptionPane.showConfirmDialog(null, "Do you want to clear data?", "Info", + int retCode = JOptionPane.showConfirmDialog(this, "Do you want to clear data?", "Info", JOptionPane.YES_NO_OPTION); String host = hostTextField.getText(); if (retCode == JOptionPane.YES_OPTION && !host.isEmpty()) { dataTabbedPane.removeAll(); splitPane.setVisible(false); + progressBar.setVisible(false); - String cleanedHost = StringProcessor.replaceFirstOccurrence(host, "*.", ""); + Config.globalDataMap.keySet().parallelStream().forEach(key -> { + if (StringProcessor.matchesHostPattern(key, host) || host.equals("*")) { + Config.globalDataMap.remove(key); + } + }); - if (host.contains("*")) { - Config.globalDataMap.keySet().removeIf(i -> i.contains(cleanedHost) || cleanedHost.contains("*")); - } else { - Config.globalDataMap.remove(host); + if (!StringProcessor.matchHostIsIp(host) && !host.contains("*.")) { + String baseDomain = StringProcessor.getBaseDomain(StringProcessor.extractHostname(host)); + + long count = Config.globalDataMap.keySet().stream() + .filter(k -> !k.equals("*." + baseDomain)) + .filter(k -> StringProcessor.matchFromEnd(k, baseDomain)) + .count(); + + if (count == 0) { + Config.globalDataMap.remove("*." + baseDomain); + } + } + + if (Config.globalDataMap.keySet().size() == 1 && Config.globalDataMap.keySet().stream().anyMatch(key -> key.contains("*"))) { + Config.globalDataMap.keySet().remove("*"); } messageTableModel.deleteByHost(host); diff --git a/src/main/java/hae/component/board/message/MessageEntry.java b/src/main/java/hae/component/board/message/MessageEntry.java index ccdf720..ec52108 100644 --- a/src/main/java/hae/component/board/message/MessageEntry.java +++ b/src/main/java/hae/component/board/message/MessageEntry.java @@ -11,8 +11,10 @@ public class MessageEntry { private final String status; private final String color; private final String method; + private final String hash; + private final String path; - MessageEntry(HttpRequestResponse requestResponse, String method, String url, String comment, String length, String color, String status) { + MessageEntry(HttpRequestResponse requestResponse, String method, String url, String comment, String length, String color, String status, String hash, String path) { this.requestResponse = requestResponse; this.method = method; this.url = url; @@ -20,6 +22,8 @@ public class MessageEntry { this.length = length; this.color = color; this.status = status; + this.hash = hash; + this.path = path; } public String getColor() { @@ -49,4 +53,12 @@ public class MessageEntry { public HttpRequestResponse getRequestResponse() { return this.requestResponse; } + + public String getHash() { + return this.hash; + } + + public String getPath() { + return this.path; + } } \ No newline at end of file diff --git a/src/main/java/hae/component/board/message/MessageTableModel.java b/src/main/java/hae/component/board/message/MessageTableModel.java index 48c6bc4..2c68cc2 100644 --- a/src/main/java/hae/component/board/message/MessageTableModel.java +++ b/src/main/java/hae/component/board/message/MessageTableModel.java @@ -11,6 +11,7 @@ import burp.api.montoya.ui.editor.HttpRequestEditor; import burp.api.montoya.ui.editor.HttpResponseEditor; import hae.Config; import hae.cache.CachePool; +import hae.utils.project.FileProcessor; import hae.utils.string.HashCalculator; import hae.utils.string.StringProcessor; @@ -22,6 +23,10 @@ import javax.swing.table.TableRowSorter; import java.nio.charset.StandardCharsets; import java.text.MessageFormat; import java.util.*; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Collectors; @@ -32,7 +37,7 @@ public class MessageTableModel extends AbstractTableModel { private final MessageTable messageTable; private final JTabbedPane messageTab; private final JSplitPane splitPane; - private final List log = new ArrayList(); + private final LinkedList log = new LinkedList<>(); private final LinkedList filteredLog; public MessageTableModel(MontoyaApi api) { @@ -92,25 +97,25 @@ public class MessageTableModel extends AbstractTableModel { splitPane.setRightComponent(messageTab); } - public void add(HttpRequestResponse messageInfo, String comment, String color) { + public void add(HttpRequestResponse messageInfo, String url, String method, String status, String length, String comment, String color, String hash, String path) { synchronized (log) { - HttpRequest httpRequest = messageInfo.request(); - String url = httpRequest.url(); - String method = httpRequest.method(); + boolean isDuplicate = false; + MessageEntry logEntry = new MessageEntry(messageInfo, method, url, comment, length, color, status, hash, path); - HttpResponse httpResponse = messageInfo.response(); - String status = String.valueOf(httpResponse.statusCode()); - String length = String.valueOf(httpResponse.body().length()); + byte[] reqByteA = new byte[0]; + byte[] resByteA = new byte[0]; - MessageEntry logEntry = new MessageEntry(messageInfo, method, url, comment, length, color, status); + if (messageInfo != null) { + HttpRequest httpRequest = messageInfo.request(); + HttpResponse httpResponse = messageInfo.response(); + reqByteA = httpRequest.toByteArray().getBytes(); + resByteA = httpResponse.toByteArray().getBytes(); + } + + // 比较Hash，如若存在重复的请求或响应，则不放入消息内容里 try { - // 比较Hash，如若存在重复的请求或响应，则不放入消息内容里 - byte[] reqByteA = httpRequest.toByteArray().getBytes(); - byte[] resByteA = httpResponse.toByteArray().getBytes(); - boolean isDuplicate = false; - - if (log.size() > 0) { + if (!log.isEmpty()) { for (MessageEntry entry : log) { HttpRequestResponse reqResMessage = entry.getRequestResponse(); byte[] reqByteB = reqResMessage.request().toByteArray().getBytes(); @@ -125,12 +130,12 @@ public class MessageTableModel extends AbstractTableModel { } } } - - if (!isDuplicate) { - log.add(logEntry); - } } catch (Exception ignored) { } + + if (!isDuplicate) { + log.add(logEntry); + } } } @@ -138,52 +143,113 @@ public class MessageTableModel extends AbstractTableModel { public void deleteByHost(String filterText) { filteredLog.clear(); List rowsToRemove = new ArrayList<>(); - for (int i = 0; i < log.size(); i++) { - MessageEntry entry = log.get(i); - String host = StringProcessor.getHostByUrl(entry.getUrl()); - if (!host.isEmpty()) { - if (StringProcessor.matchesHostPattern(host, filterText) || filterText.contains("*")) { - rowsToRemove.add(i); + + new SwingWorker() { + @Override + protected Void doInBackground() { + for (int i = 0; i < log.size(); i++) { + MessageEntry entry = log.get(i); + String host = StringProcessor.getHostByUrl(entry.getUrl()); + if (!host.isEmpty()) { + if (StringProcessor.matchesHostPattern(host, filterText) || filterText.equals("*")) { + rowsToRemove.add(i); + } + } + } + + for (int i = rowsToRemove.size() - 1; i >= 0; i--) { + int row = rowsToRemove.get(i); + log.remove(row); + } + + return null; + } + + @Override + protected void done() { + if (!rowsToRemove.isEmpty()) { + int[] rows = rowsToRemove.stream().mapToInt(Integer::intValue).toArray(); + fireTableRowsDeleted(rows[0], rows[rows.length - 1]); } } - } - - for (int i = rowsToRemove.size() - 1; i >= 0; i--) { - int row = rowsToRemove.get(i); - log.remove(row); - } - - if (!rowsToRemove.isEmpty()) { - int[] rows = rowsToRemove.stream().mapToInt(Integer::intValue).toArray(); - fireTableRowsDeleted(rows[0], rows[rows.length - 1]); - } + }.execute(); } public void applyHostFilter(String filterText) { filteredLog.clear(); - fireTableDataChanged(); - String cleanedText = StringProcessor.replaceFirstOccurrence(filterText, "*.", ""); - for (MessageEntry entry : log) { - String host = StringProcessor.getHostByUrl(entry.getUrl()); - if (!host.isEmpty()) { - if (filterText.contains("*.") && StringProcessor.matchFromEnd(StringProcessor.extractHostname(host), cleanedText)) { - filteredLog.add(entry); - } else if (host.equals(filterText) || filterText.contains("*")) { - filteredLog.add(entry); + ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2); + List> futures = new ArrayList<>(); + try { + log.parallelStream().forEach(entry -> { + Future future = executor.submit(() -> { + MessageEntry finalEntry = getEntryByFile(entry); + String host = StringProcessor.getHostByUrl(finalEntry.getUrl()); + if (!host.isEmpty()) { + synchronized (filteredLog) { + if (StringProcessor.matchesHostPattern(host, filterText) || filterText.contains("*")) { + filteredLog.add(finalEntry); + } + } + } + }); + + futures.add(future); + }); + + for (Future future : futures) { + try { + future.get(); + } catch (InterruptedException | ExecutionException e) { + Thread.currentThread().interrupt(); } } + + } catch (Exception e) { + api.logging().logToError("applyHostFilter: " + e.getMessage()); + } finally { + executor.shutdown(); + fireTableDataChanged(); + } + } + + private MessageEntry getEntryByFile(MessageEntry entry) { + HttpRequestResponse requestResponse = entry.getRequestResponse(); + if (requestResponse == null) { + String url = entry.getUrl(); + String method = entry.getMethod(); + String status = entry.getStatus(); + String comment = entry.getComment(); + String color = entry.getColor(); + String path = entry.getPath(); + String hash = entry.getHash(); + int length = Integer.parseInt(entry.getLength()); + + byte[] contents = FileProcessor.readFileContent(path, hash); + + if (contents.length > length) { + byte[] response = Arrays.copyOf(contents, length); + byte[] request = Arrays.copyOfRange(contents, length, contents.length); + requestResponse = StringProcessor.createHttpRequestResponse(url, request, response); + + int index = log.indexOf(entry); + entry = new MessageEntry(requestResponse, method, url, comment, String.valueOf(length), color, status, "", ""); + log.set(index, entry); + } } - fireTableDataChanged(); + return entry; } public void applyMessageFilter(String tableName, String filterText) { filteredLog.clear(); for (MessageEntry entry : log) { HttpRequestResponse requestResponse = entry.getRequestResponse(); - HttpRequest httpRequest = requestResponse.request(); - HttpResponse httpResponse = requestResponse.response(); + // 标志变量，表示是否满足过滤条件 + AtomicBoolean isMatched = new AtomicBoolean(false); + + HttpRequest httpRequest = entry.getRequestResponse().request(); + HttpResponse httpResponse = entry.getRequestResponse().response(); String requestString = new String(httpRequest.toByteArray().getBytes(), StandardCharsets.UTF_8); String requestBody = new String(httpRequest.body().getBytes(), StandardCharsets.UTF_8); @@ -197,9 +263,7 @@ public class MessageTableModel extends AbstractTableModel { .map(HttpHeader::toString) .collect(Collectors.joining("\n")); - // 标志变量，表示是否满足过滤条件 - AtomicBoolean isMatched = new AtomicBoolean(false); - + MessageEntry finalEntry = entry; Config.globalRules.keySet().forEach(i -> { for (Object[] objects : Config.globalRules.get(i)) { String name = objects[1].toString(); @@ -207,7 +271,7 @@ public class MessageTableModel extends AbstractTableModel { String scope = objects[6].toString(); // 从注释中查看是否包含当前规则名，包含的再进行查询，有效减少无意义的检索时间 - if (entry.getComment().contains(name)) { + if (finalEntry.getComment().contains(name)) { if (name.equals(tableName)) { // 标志变量，表示当前规则是否匹配 boolean isMatch = false; @@ -305,7 +369,7 @@ public class MessageTableModel extends AbstractTableModel { if (!map2.containsKey(key)) { return false; } - if (!areInnerMapsEqual(map1.get(key), map2.get(key))) { + if (areInnerMapsEqual(map1.get(key), map2.get(key))) { return false; } } @@ -315,30 +379,29 @@ public class MessageTableModel extends AbstractTableModel { private boolean areInnerMapsEqual(Map innerMap1, Map innerMap2) { if (innerMap1.size() != innerMap2.size()) { - return false; + return true; } for (String key : innerMap1.keySet()) { if (!innerMap2.containsKey(key)) { - return false; + return true; } Object value1 = innerMap1.get(key); Object value2 = innerMap2.get(key); // 如果值是Map，则递归对比 if (value1 instanceof Map && value2 instanceof Map) { - if (!areInnerMapsEqual((Map) value1, (Map) value2)) { - return false; + if (areInnerMapsEqual((Map) value1, (Map) value2)) { + return true; } } else if (!value1.equals(value2)) { - return false; + return true; } } - return true; + return false; } - public JSplitPane getSplitPane() { return splitPane; } @@ -366,6 +429,7 @@ public class MessageTableModel extends AbstractTableModel { if (filteredLog.isEmpty()) { return ""; } + MessageEntry messageEntry = filteredLog.get(rowIndex); return switch (columnIndex) { @@ -393,7 +457,7 @@ public class MessageTableModel extends AbstractTableModel { } public class MessageTable extends JTable { - private MessageEntry MessageEntry; + private MessageEntry messageEntry; private SwingWorker currentWorker; // 设置响应报文返回的最大长度 private final int MAX_LENGTH = 5242880; @@ -413,7 +477,7 @@ public class MessageTableModel extends AbstractTableModel { int selectedIndex = convertRowIndexToModel(row); if (lastSelectedIndex != selectedIndex) { lastSelectedIndex = selectedIndex; - MessageEntry = filteredLog.get(selectedIndex); + messageEntry = filteredLog.get(selectedIndex); requestEditor.setRequest(HttpRequest.httpRequest("Loading...")); responseEditor.setResponse(HttpResponse.httpResponse("Loading...")); @@ -425,8 +489,10 @@ public class MessageTableModel extends AbstractTableModel { currentWorker = new SwingWorker<>() { @Override protected ByteArray[] doInBackground() { - ByteArray requestByte = MessageEntry.getRequestResponse().request().toByteArray(); - ByteArray responseByte = MessageEntry.getRequestResponse().response().toByteArray(); + HttpRequestResponse httpRequestResponse = messageEntry.getRequestResponse(); + + ByteArray requestByte = messageEntry.getRequestResponse().request().toByteArray(); + ByteArray responseByte = messageEntry.getRequestResponse().response().toByteArray(); if (responseByte.length() > MAX_LENGTH) { String ellipsis = "\r\n......"; @@ -441,7 +507,7 @@ public class MessageTableModel extends AbstractTableModel { if (!isCancelled()) { try { ByteArray[] result = (ByteArray[]) get(); - requestEditor.setRequest(HttpRequest.httpRequest(MessageEntry.getRequestResponse().httpService(), result[0])); + requestEditor.setRequest(HttpRequest.httpRequest(messageEntry.getRequestResponse().httpService(), result[0])); responseEditor.setResponse(HttpResponse.httpResponse(result[1])); } catch (Exception ignored) { } diff --git a/src/main/java/hae/component/config/Config.java b/src/main/java/hae/component/config/Config.java index 7375c83..e7cfc63 100644 --- a/src/main/java/hae/component/config/Config.java +++ b/src/main/java/hae/component/config/Config.java @@ -250,7 +250,7 @@ public class Config extends JPanel { private void onlineUpdateActionPerformed(ActionEvent e) { // 添加提示框防止用户误触导致配置更新 - int retCode = JOptionPane.showConfirmDialog(null, "Do you want to update rules?", "Info", JOptionPane.YES_NO_OPTION); + int retCode = JOptionPane.showConfirmDialog(this, "Do you want to update rules?", "Info", JOptionPane.YES_NO_OPTION); if (retCode == JOptionPane.YES_OPTION) { configLoader.initRulesByNet(); reloadActionPerformed(null); diff --git a/src/main/java/hae/component/rule/Rule.java b/src/main/java/hae/component/rule/Rule.java index 72143b3..49078a7 100644 --- a/src/main/java/hae/component/rule/Rule.java +++ b/src/main/java/hae/component/rule/Rule.java @@ -97,7 +97,7 @@ public class Rule extends JPanel { Display ruleDisplay = new Display(); ruleDisplay.formatTextField.setText("{0}"); - int showState = JOptionPane.showConfirmDialog(null, ruleDisplay, "Add Rule", JOptionPane.OK_OPTION); + int showState = JOptionPane.showConfirmDialog(this, ruleDisplay, "Add Rule", JOptionPane.OK_OPTION); if (showState == YES_OPTION) { Vector ruleData = new Vector<>(); ruleData.add(false); @@ -132,7 +132,7 @@ public class Rule extends JPanel { ruleDisplay.formatTextField.setEnabled(ruleDisplay.engineComboBox.getSelectedItem().toString().equals("nfa")); - int showState = JOptionPane.showConfirmDialog(null, ruleDisplay, "Edit Rule", JOptionPane.OK_OPTION); + int showState = JOptionPane.showConfirmDialog(this, ruleDisplay, "Edit Rule", JOptionPane.OK_OPTION); if (showState == 0) { int select = ruleTable.convertRowIndexToModel(ruleTable.getSelectedRow()); model.setValueAt(ruleDisplay.ruleNameTextField.getText(), select, 1); @@ -151,7 +151,7 @@ public class Rule extends JPanel { private void ruleRemoveActionPerformed(ActionEvent e, JTable ruleTable, JTabbedPane tabbedPane) { if (ruleTable.getSelectedRowCount() >= 1) { - if (JOptionPane.showConfirmDialog(null, "Are you sure you want to remove this rule?", "Info", JOptionPane.YES_NO_OPTION) == 0) { + if (JOptionPane.showConfirmDialog(this, "Are you sure you want to remove this rule?", "Info", JOptionPane.YES_NO_OPTION) == 0) { DefaultTableModel model = (DefaultTableModel) ruleTable.getModel(); int select = ruleTable.convertRowIndexToModel(ruleTable.getSelectedRow()); diff --git a/src/main/java/hae/component/rule/Rules.java b/src/main/java/hae/component/rule/Rules.java index 1bc87f7..0df6113 100644 --- a/src/main/java/hae/component/rule/Rules.java +++ b/src/main/java/hae/component/rule/Rules.java @@ -109,7 +109,7 @@ public class Rules extends JTabbedPane { private void deleteRuleGroupActionPerformed(ActionEvent e) { if (getTabCount() > 2) { - int retCode = JOptionPane.showConfirmDialog(null, "Do you want to delete this rule group?", "Info", + int retCode = JOptionPane.showConfirmDialog(this, "Do you want to delete this rule group?", "Info", JOptionPane.YES_NO_OPTION); if (retCode == JOptionPane.YES_OPTION) { String title = getTitleAt(getSelectedIndex()); diff --git a/src/main/java/hae/instances/editor/RequestEditor.java b/src/main/java/hae/instances/editor/RequestEditor.java index a1d8325..e71074e 100644 --- a/src/main/java/hae/instances/editor/RequestEditor.java +++ b/src/main/java/hae/instances/editor/RequestEditor.java @@ -122,7 +122,9 @@ public class RequestEditor implements HttpRequestEditorProvider { boolean isBlockHost = false; for (String hostName : hostList) { String cleanedHost = StringProcessor.replaceFirstOccurrence(hostName, "*.", ""); - if (StringProcessor.matchFromEnd(host, cleanedHost)) { + if (hostName.contains("*.") && StringProcessor.matchFromEnd(host, cleanedHost)) { + isBlockHost = true; + } else if (host.equals(hostName) || hostName.equals("*")) { isBlockHost = true; } } diff --git a/src/main/java/hae/instances/http/HttpMessageHandler.java b/src/main/java/hae/instances/http/HttpMessageHandler.java index 8610e4f..f5f63fe 100644 --- a/src/main/java/hae/instances/http/HttpMessageHandler.java +++ b/src/main/java/hae/instances/http/HttpMessageHandler.java @@ -80,7 +80,12 @@ public class HttpMessageHandler implements HttpHandler { HttpRequestResponse httpRequestResponse = HttpRequestResponse.httpRequestResponse(httpRequest.get(), httpResponseReceived); // 添加到Databoard - messageTableModel.add(httpRequestResponse, comment, color); + String method = httpRequest.get().method(); + String url = httpRequest.get().url(); + String status = String.valueOf(httpResponseReceived.statusCode()); + String length = String.valueOf(httpResponseReceived.toByteArray().length()); + + messageTableModel.add(httpRequestResponse, url, method, status, length, comment, color, "", ""); } } diff --git a/src/main/java/hae/instances/http/utils/RegularMatcher.java b/src/main/java/hae/instances/http/utils/RegularMatcher.java index b92b8b6..6a01d61 100644 --- a/src/main/java/hae/instances/http/utils/RegularMatcher.java +++ b/src/main/java/hae/instances/http/utils/RegularMatcher.java @@ -126,7 +126,7 @@ public class RegularMatcher { String[] splitHost = host.split("\\."); String onlyHost = host.split(":")[0]; - String anyHost = (splitHost.length > 2 && !onlyHost.matches("\\b(?:\\d{1,3}\\.){3}\\d{1,3}\\b")) ? StringProcessor.replaceFirstOccurrence(onlyHost, splitHost[0], "*") : ""; + String anyHost = (splitHost.length > 2 && !StringProcessor.matchHostIsIp(onlyHost)) ? StringProcessor.replaceFirstOccurrence(onlyHost, splitHost[0], "*") : ""; if (!Config.globalDataMap.containsKey(anyHost) && anyHost.length() > 0) { // 添加通配符Host，实际数据从查询哪里将所有数据提取 diff --git a/src/main/java/hae/utils/project/FileProcessor.java b/src/main/java/hae/utils/project/FileProcessor.java new file mode 100644 index 0000000..7e37b6d --- /dev/null +++ b/src/main/java/hae/utils/project/FileProcessor.java @@ -0,0 +1,47 @@ +package hae.utils.project; + +import java.io.File; +import java.nio.file.DirectoryStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Comparator; + +public class FileProcessor { + public static void deleteDirectoryWithContents(Path pathToBeDeleted) { + if (pathToBeDeleted != null) { + try { + Files.walk(pathToBeDeleted) + .sorted(Comparator.reverseOrder()) + .map(Path::toFile) + .forEach(File::delete); + } catch (Exception ignored) { + } + } + } + + public static byte[] readFileContent(String basePath, String fileName) { + Path filePath = Paths.get(basePath, fileName); + Path path = Paths.get(basePath); + try { + byte[] fileContent = Files.readAllBytes(filePath); + + Files.deleteIfExists(filePath); + + boolean isEmpty = isDirectoryEmpty(path); + if (isEmpty) { + Files.deleteIfExists(path); + } + + return fileContent; + } catch (Exception e) { + return new byte[0]; + } + } + + private static boolean isDirectoryEmpty(Path directory) throws Exception { + try (DirectoryStream dirStream = Files.newDirectoryStream(directory)) { + return !dirStream.iterator().hasNext(); + } + } +} diff --git a/src/main/java/hae/utils/project/ProjectProcessor.java b/src/main/java/hae/utils/project/ProjectProcessor.java index 064e107..7d1a2d6 100644 --- a/src/main/java/hae/utils/project/ProjectProcessor.java +++ b/src/main/java/hae/utils/project/ProjectProcessor.java @@ -2,14 +2,21 @@ package hae.utils.project; import burp.api.montoya.MontoyaApi; import hae.utils.project.model.HaeFileContent; +import org.yaml.snakeyaml.LoaderOptions; import org.yaml.snakeyaml.Yaml; import java.io.*; import java.nio.charset.StandardCharsets; -import java.util.List; -import java.util.Map; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardCopyOption; +import java.util.*; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.zip.ZipEntry; -import java.util.zip.ZipInputStream; +import java.util.zip.ZipFile; import java.util.zip.ZipOutputStream; public class ProjectProcessor { @@ -19,59 +26,169 @@ public class ProjectProcessor { this.api = api; } - public boolean createHaeFile(String haeFilePath, String host, Map> dataMap, Map> httpMap) { + public boolean createHaeFile(String haeFilePath, String host, Map> dataMap, Map> urlMap, Map> httpMap) { + ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors()); + List> futures = new ArrayList<>(); + ByteArrayOutputStream dataYamlStream = new ByteArrayOutputStream(); - ByteArrayOutputStream httpYamlStream = new ByteArrayOutputStream(); + ByteArrayOutputStream urlYamlStream = new ByteArrayOutputStream(); Yaml yaml = new Yaml(); yaml.dump(dataMap, new OutputStreamWriter(dataYamlStream, StandardCharsets.UTF_8)); - yaml.dump(httpMap, new OutputStreamWriter(httpYamlStream, StandardCharsets.UTF_8)); + yaml.dump(urlMap, new OutputStreamWriter(urlYamlStream, StandardCharsets.UTF_8)); try (ZipOutputStream zipOut = new ZipOutputStream(new FileOutputStream(haeFilePath))) { zipOut.putNextEntry(new ZipEntry("info")); zipOut.write(host.getBytes(StandardCharsets.UTF_8)); zipOut.closeEntry(); - zipOut.putNextEntry(new ZipEntry("data.yml")); + zipOut.putNextEntry(new ZipEntry("data")); zipOut.write(dataYamlStream.toByteArray()); zipOut.closeEntry(); - zipOut.putNextEntry(new ZipEntry("http.yml")); - zipOut.write(httpYamlStream.toByteArray()); + zipOut.putNextEntry(new ZipEntry("url")); + zipOut.write(urlYamlStream.toByteArray()); zipOut.closeEntry(); + + for (String httpHash : httpMap.keySet()) { + Map httpItem = httpMap.get(httpHash); + futures.add(executorService.submit(() -> { + try { + ByteArrayOutputStream httpOutStream = new ByteArrayOutputStream(); + byte[] request = (byte[]) httpItem.get("request"); + byte[] response = (byte[]) httpItem.get("response"); + + httpOutStream.write(response); + httpOutStream.write(request); + + synchronized (zipOut) { + zipOut.putNextEntry(new ZipEntry(String.format("http/%s", httpHash))); + zipOut.write(httpOutStream.toByteArray()); + zipOut.closeEntry(); + } + } catch (Exception e) { + api.logging().logToError("createHaeFile: " + e.getMessage()); + } + })); + } + + for (Future future : futures) { + try { + future.get(); + } catch (InterruptedException | ExecutionException e) { + Thread.currentThread().interrupt(); + } + } } catch (Exception e) { - api.logging().logToOutput(e.getMessage()); + api.logging().logToError("createHaeFile: " + e.getMessage()); return false; + } finally { + executorService.shutdown(); } return true; } public HaeFileContent readHaeFile(String haeFilePath) { - HaeFileContent haeFileContent = new HaeFileContent(api); - Yaml yaml = new Yaml(); + ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors()); + List> futures = new ArrayList<>(); - try (ZipInputStream zipIn = new ZipInputStream(new FileInputStream(haeFilePath))) { - ZipEntry entry; - while ((entry = zipIn.getNextEntry()) != null) { - switch (entry.getName()) { - case "info": - haeFileContent.setHost(new String(zipIn.readAllBytes(), StandardCharsets.UTF_8)); - break; - case "data.yml": - haeFileContent.setDataMap(yaml.load(new InputStreamReader(zipIn, StandardCharsets.UTF_8))); - break; - case "http.yml": - haeFileContent.setHttpMap(yaml.load(new InputStreamReader(zipIn, StandardCharsets.UTF_8))); - break; + HaeFileContent haeFileContent = new HaeFileContent(api); // 假设api是正确的 + LoaderOptions loaderOptions = new LoaderOptions(); + loaderOptions.setMaxAliasesForCollections(Integer.MAX_VALUE); + loaderOptions.setCodePointLimit(Integer.MAX_VALUE); + Yaml yaml = new Yaml(loaderOptions); + Path tempDirectory = null; + + try { + if (hasValidStructure(haeFilePath)) { + tempDirectory = Files.createTempDirectory("hae"); + haeFileContent.setHttpPath(tempDirectory.toString()); + + try (ZipFile zipFile = new ZipFile(haeFilePath)) { + Enumeration entries = zipFile.entries(); + while (entries.hasMoreElements()) { + ZipEntry entry = entries.nextElement(); + String fileName = entry.getName(); + if (fileName.startsWith("http/")) { + Path filePath = tempDirectory.resolve(fileName.substring("http/".length())); + futures.add(executorService.submit(() -> { + try (InputStream in = zipFile.getInputStream(entry)) { + Files.copy(in, filePath, StandardCopyOption.REPLACE_EXISTING); + } catch (IOException e) { + api.logging().logToError("readHaeFile: " + e.getMessage()); + } + })); + } else { + try (InputStream in = zipFile.getInputStream(entry)) { + switch (fileName) { + case "info" -> + haeFileContent.setHost(new String(in.readAllBytes(), StandardCharsets.UTF_8)); + case "data" -> + haeFileContent.setDataMap(yaml.load(new InputStreamReader(in, StandardCharsets.UTF_8))); + case "url" -> + haeFileContent.setUrlMap(yaml.load(new InputStreamReader(in, StandardCharsets.UTF_8))); + } + } + } + } + + for (Future future : futures) { + try { + future.get(); + } catch (InterruptedException | ExecutionException e) { + Thread.currentThread().interrupt(); + } + } } - zipIn.closeEntry(); } } catch (Exception e) { - api.logging().logToOutput(e.getMessage()); - return null; + api.logging().logToError("readHaeFile: " + e.getMessage()); + if (tempDirectory != null) { + FileProcessor.deleteDirectoryWithContents(tempDirectory); + } + haeFileContent = null; + } finally { + executorService.shutdown(); } + return haeFileContent; } + + private boolean hasValidStructure(String zipFilePath) { + Set requiredRootEntries = new HashSet<>(); + requiredRootEntries.add("info"); + requiredRootEntries.add("data"); + requiredRootEntries.add("url"); + + boolean hasHttpDirectoryWithFiles = false; + + try { + ZipFile zipFile = new ZipFile(zipFilePath); + Enumeration entries = zipFile.entries(); + + while (entries.hasMoreElements()) { + ZipEntry entry = entries.nextElement(); + String name = entry.getName(); + + if (!entry.isDirectory() && !name.contains("/")) { + requiredRootEntries.remove(name); + } + + if (name.startsWith("http/") && !entry.isDirectory()) { + hasHttpDirectoryWithFiles = true; + } + + if (requiredRootEntries.isEmpty() && hasHttpDirectoryWithFiles) { + break; + } + } + + zipFile.close(); + } catch (Exception ignored) { + } + + return requiredRootEntries.isEmpty() && hasHttpDirectoryWithFiles; + } } diff --git a/src/main/java/hae/utils/project/model/HaeFileContent.java b/src/main/java/hae/utils/project/model/HaeFileContent.java index 60e8eaa..5782d70 100644 --- a/src/main/java/hae/utils/project/model/HaeFileContent.java +++ b/src/main/java/hae/utils/project/model/HaeFileContent.java @@ -11,13 +11,14 @@ import java.util.Map; public class HaeFileContent { private final MontoyaApi api; private String host; + private String httpPath; private final Map> dataMap; - private final Map> httpMap; + private final Map> urlMap; public HaeFileContent(MontoyaApi api) { this.api = api; this.dataMap = new HashMap<>(); - this.httpMap = new HashMap<>(); + this.urlMap = new HashMap<>(); } public String getHost() { @@ -28,14 +29,22 @@ public class HaeFileContent { return dataMap; } - public Map> getHttpMap() { - return httpMap; + public Map> getUrlMap() { + return urlMap; + } + + public String getHttpPath() { + return httpPath; } public void setHost(String host) { this.host = host; } + public void setHttpPath(String path) { + this.httpPath = path; + } + public void setDataMap(Map> dataMap) { for (Map.Entry> entry : dataMap.entrySet()) { List values = new ArrayList<>(); @@ -50,8 +59,8 @@ public class HaeFileContent { } } - public void setHttpMap(Map> httpMap) { - for (Map.Entry> entry : httpMap.entrySet()) { + public void setUrlMap(Map> urlMap) { + for (Map.Entry> entry : urlMap.entrySet()) { Map newValues = new HashMap<>(); Map values = entry.getValue(); for (String key : values.keySet()) { @@ -61,7 +70,7 @@ public class HaeFileContent { newValues.put(key, values.get(key).toString()); } } - this.httpMap.put(entry.getKey(), newValues); + this.urlMap.put(entry.getKey(), newValues); } } } \ No newline at end of file diff --git a/src/main/java/hae/utils/string/StringProcessor.java b/src/main/java/hae/utils/string/StringProcessor.java index 7e9644b..7168b75 100644 --- a/src/main/java/hae/utils/string/StringProcessor.java +++ b/src/main/java/hae/utils/string/StringProcessor.java @@ -1,8 +1,17 @@ package hae.utils.string; +import burp.api.montoya.core.ByteArray; +import burp.api.montoya.http.HttpService; +import burp.api.montoya.http.message.HttpRequestResponse; +import burp.api.montoya.http.message.requests.HttpRequest; +import burp.api.montoya.http.message.responses.HttpResponse; + import java.net.URL; +import java.time.LocalDateTime; +import java.time.format.DateTimeFormatter; import java.util.HashMap; import java.util.Map; +import java.util.UUID; public class StringProcessor { public static String replaceFirstOccurrence(String original, String find, String replace) { @@ -55,6 +64,24 @@ public class StringProcessor { return matchesDirectly || matchesPattern; } + public static HttpRequestResponse createHttpRequestResponse(String url, byte[] request, byte[] response) { + HttpService httpService = HttpService.httpService(url); + HttpRequest httpRequest = HttpRequest.httpRequest(httpService, ByteArray.byteArray(request)); + HttpResponse httpResponse = HttpResponse.httpResponse(ByteArray.byteArray(response)); + return HttpRequestResponse.httpRequestResponse(httpRequest, httpResponse); + } + + public static String getCurrentTime() { + DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyyMMdd_HHmmss"); + LocalDateTime now = LocalDateTime.now(); + return now.format(formatter); + } + + public static String getRandomUUID() { + UUID uuid = UUID.randomUUID(); + return uuid.toString(); + } + public static String mergeComment(String comment) { if (!comment.contains(",")) { return comment; @@ -92,6 +119,21 @@ public class StringProcessor { return host; } + public static String getBaseDomain(String host) { + int lastIndex = host.lastIndexOf('.'); + if (lastIndex > 0) { + int secondLastIndex = host.substring(0, lastIndex).lastIndexOf('.'); + if (secondLastIndex >= 0) { + return host.substring(secondLastIndex + 1); + } + } + return host; + } + + public static boolean matchHostIsIp(String host) { + return host.matches("\\b(?:\\d{1,3}\\.){3}\\d{1,3}\\b"); + } + private static Map getStringIntegerMap(String comment) { Map itemCounts = new HashMap<>(); String[] items = comment.split(", ");