From e5f55b6c4c05f4e051c535a8c5d6a20b1a83329e Mon Sep 17 00:00:00 2001 From: gh0stkey <24655118+gh0stkey@users.noreply.github.com> Date: Wed, 19 Jun 2024 22:16:57 +0800 Subject: [PATCH] Version: 3.2.2 Update --- src/main/java/hae/HaE.java | 13 +- .../java/hae/component/board/Databoard.java | 151 +++++++++++------- .../java/hae/component/board/Datatable.java | 8 +- .../board/message/MessageRenderer.java | 6 +- .../board/message/MessageTableModel.java | 114 +++++-------- .../instances/http/HttpMessageHandler.java | 34 ++-- .../http/utils/MessageProcessor.java | 2 +- .../hae/utils/project/ProjectProcessor.java | 45 +++--- .../hae/utils/string/StringProcessor.java | 11 -- 9 files changed, 192 insertions(+), 192 deletions(-) diff --git a/src/main/java/hae/HaE.java b/src/main/java/hae/HaE.java index c0f408b..6e23f85 100644 --- a/src/main/java/hae/HaE.java +++ b/src/main/java/hae/HaE.java @@ -2,7 +2,9 @@ import burp.api.montoya.BurpExtension; import burp.api.montoya.MontoyaApi; +import burp.api.montoya.extension.ExtensionUnloadingHandler; import burp.api.montoya.logging.Logging; +import hae.cache.CachePool; import hae.component.Main; import hae.component.board.message.MessageTableModel; import hae.instances.editor.RequestEditor; @@ -16,7 +18,7 @@ public class HaE implements BurpExtension { @Override public void initialize(MontoyaApi api) { // 设置扩展名称 - String version = "3.2.1"; + String version = "3.2.2"; api.extension().setName(String.format("HaE (%s) - Highlighter and Extractor", version)); // 加载扩展后输出的项目信息 @@ -43,5 +45,14 @@ public void initialize(MontoyaApi api) { api.userInterface().registerHttpRequestEditorProvider(new RequestEditor(api, configLoader)); api.userInterface().registerHttpResponseEditorProvider(new ResponseEditor(api, configLoader)); api.userInterface().registerWebSocketMessageEditorProvider(new WebSocketEditor(api)); + + api.extension().registerUnloadingHandler(new ExtensionUnloadingHandler() { + @Override + public void extensionUnloaded() { + // 卸载清空数据 + Config.globalDataMap.clear(); + CachePool.clear(); + } + }); } } diff --git a/src/main/java/hae/component/board/Databoard.java b/src/main/java/hae/component/board/Databoard.java index f491a14..e71cf98 100644 --- a/src/main/java/hae/component/board/Databoard.java +++ b/src/main/java/hae/component/board/Databoard.java @@ -23,7 +23,10 @@ import java.io.File; import java.util.List; import java.util.*; -import java.util.concurrent.*; +import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import java.util.function.Function; import java.util.stream.Collectors; @@ -43,7 +46,10 @@ public class Databoard extends JPanel { private final DefaultComboBoxModel comboBoxModel = new DefaultComboBoxModel(); private final JComboBox hostComboBox = new JComboBox(comboBoxModel); - private SwingWorker currentWorker; + private SwingWorker handleComboBoxWorker; + private SwingWorker applyHostFilterWorker; + private SwingWorker, Void> exportActionWorker; + private SwingWorker, Void> importActionWorker; public Databoard(MontoyaApi api, ConfigLoader configLoader, MessageTableModel messageTableModel) { this.api = api; @@ -77,11 +83,10 @@ private void initComponents() { hostTextField = new JTextField(); splitPane = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT); - dataTabbedPane = new JTabbedPane(JTabbedPane.TOP); + 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; @@ -185,42 +190,46 @@ private void handleComboBoxAction(ActionEvent e) { progressBar.setVisible(true); setProgressBar(true); String selectedHost = hostComboBox.getSelectedItem().toString(); - hostTextField.setText(selectedHost); - if (currentWorker != null && !currentWorker.isDone()) { - currentWorker.cancel(true); - } + if (getHostByList().contains(selectedHost)) { + hostTextField.setText(selectedHost); - currentWorker = new SwingWorker() { - @Override - protected Boolean doInBackground() { - return populateTabbedPaneByHost(selectedHost); + if (handleComboBoxWorker != null && !handleComboBoxWorker.isDone()) { + handleComboBoxWorker.cancel(true); } - @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); + handleComboBoxWorker = 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.setLeftComponent(dataTabbedPane); + splitPane.setRightComponent(messageSplitPane); + messageTable = messageTableModel.getMessageTable(); + resizePanel(); + + splitPane.setVisible(true); + hostTextField.setText(selectedHost); + + hostComboBox.setPopupVisible(false); + applyHostFilter(selectedHost); + } + } catch (Exception ignored) { } - } catch (Exception ignored) { } } - } - }; + }; - currentWorker.execute(); + handleComboBoxWorker.execute(); + } } } @@ -254,7 +263,7 @@ private boolean populateTabbedPaneByHost(String selectedHost) { if (selectedHost.contains("*")) { selectedDataMap = new HashMap<>(); - dataMap.keySet().parallelStream().forEach(key -> { + dataMap.keySet().forEach(key -> { if ((StringProcessor.matchesHostPattern(key, selectedHost) || selectedHost.equals("*")) && !key.contains("*")) { Map> ruleMap = dataMap.get(key); for (String ruleKey : ruleMap.keySet()) { @@ -286,6 +295,7 @@ private boolean populateTabbedPaneByHost(String selectedHost) { return true; } + return false; } @@ -316,7 +326,11 @@ private void applyHostFilter(String filterText) { TableRowSorter sorter = (TableRowSorter) messageTable.getRowSorter(); String cleanedText = StringProcessor.replaceFirstOccurrence(filterText, "*.", ""); - new SwingWorker() { + if (applyHostFilterWorker != null && !applyHostFilterWorker.isDone()) { + applyHostFilterWorker.cancel(true); + } + + applyHostFilterWorker = new SwingWorker() { @Override protected Void doInBackground() throws Exception { RowFilter rowFilter = new RowFilter() { @@ -340,8 +354,9 @@ public boolean include(Entry entry) { protected void done() { setProgressBar(false); } - }.execute(); + }; + applyHostFilterWorker.execute(); } private List getHostByList() { @@ -364,7 +379,11 @@ private void exportActionPerformed(ActionEvent e) { return; } - new SwingWorker, Void>() { + if (exportActionWorker != null && !exportActionWorker.isDone()) { + exportActionWorker.cancel(true); + } + + exportActionWorker = new SwingWorker, Void>() { @Override protected List doInBackground() { ConcurrentHashMap>> dataMap = Config.globalDataMap; @@ -382,7 +401,9 @@ protected void done() { } catch (Exception ignored) { } } - }.execute(); + }; + + exportActionWorker.execute(); } private List exportData(String selectedHost, String exportDir, Map>> dataMap) { @@ -469,7 +490,11 @@ private void importActionPerformed(ActionEvent e) { return; } - new SwingWorker, Void>() { + if (importActionWorker != null && !importActionWorker.isDone()) { + importActionWorker.cancel(true); + } + + importActionWorker = new SwingWorker, Void>() { @Override protected List doInBackground() { List filesWithExtension = findFilesWithExtension(new File(exportDir), ".hae"); @@ -489,15 +514,18 @@ protected void done() { } catch (Exception ignored) { } } - }.execute(); + }; + + importActionWorker.execute(); } private String importData(String filename) { + ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2); + HaeFileContent haeFileContent = projectProcessor.readHaeFile(filename); boolean readStatus = haeFileContent != null; - ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2); - List> futures = new ArrayList<>(); + List> tasks = new ArrayList<>(); if (readStatus) { try { @@ -505,7 +533,7 @@ private String importData(String filename) { haeFileContent.getDataMap().forEach((key, value) -> RegularMatcher.putDataToGlobalMap(host, key, value)); haeFileContent.getUrlMap().forEach((key, urlItemMap) -> { - Future future = executor.submit(() -> { + tasks.add(() -> { String url = urlItemMap.get("url"); String comment = urlItemMap.get("comment"); String color = urlItemMap.get("color"); @@ -515,18 +543,11 @@ private String importData(String filename) { String path = haeFileContent.getHttpPath(); messageTableModel.add(null, url, method, status, length, comment, color, key, path); + return null; }); - - futures.add(future); }); - for (Future future : futures) { - try { - future.get(); - } catch (InterruptedException | ExecutionException e) { - Thread.currentThread().interrupt(); - } - } + executor.invokeAll(tasks); } catch (Exception e) { api.logging().logToError("importData: " + e.getMessage()); } finally { @@ -590,20 +611,26 @@ private void clearActionPerformed(ActionEvent e) { } }); - if (!StringProcessor.matchHostIsIp(host) && !host.contains("*.")) { - String baseDomain = StringProcessor.getBaseDomain(StringProcessor.extractHostname(host)); + // 删除无用的数据 + Set wildcardKeys = Config.globalDataMap.keySet().stream() + .filter(key -> key.startsWith("*.")) + .collect(Collectors.toSet()); - long count = Config.globalDataMap.keySet().stream() - .filter(k -> !k.equals("*." + baseDomain)) - .filter(k -> StringProcessor.matchFromEnd(k, baseDomain)) - .count(); + Set existingSuffixes = Config.globalDataMap.keySet().stream() + .filter(key -> !key.startsWith("*.")) + .map(key -> { + int dotIndex = key.indexOf("."); + return dotIndex != -1 ? key.substring(dotIndex) : ""; + }) + .collect(Collectors.toSet()); - if (count == 0) { - Config.globalDataMap.remove("*." + baseDomain); - } - } + Set keysToRemove = wildcardKeys.stream() + .filter(key -> !existingSuffixes.contains(key.substring(1))) + .collect(Collectors.toSet()); + + keysToRemove.forEach(Config.globalDataMap::remove); - if (Config.globalDataMap.keySet().size() == 1 && Config.globalDataMap.keySet().stream().anyMatch(key -> key.contains("*"))) { + if (Config.globalDataMap.keySet().size() == 1 && Config.globalDataMap.keySet().stream().anyMatch(key -> key.equals("*"))) { Config.globalDataMap.keySet().remove("*"); } diff --git a/src/main/java/hae/component/board/Datatable.java b/src/main/java/hae/component/board/Datatable.java index d0fccee..2626ba5 100644 --- a/src/main/java/hae/component/board/Datatable.java +++ b/src/main/java/hae/component/board/Datatable.java @@ -158,7 +158,7 @@ public void setTableListener(MessageTableModel messagePanel) { @Override public void exportToClipboard(JComponent comp, Clipboard clip, int action) throws IllegalStateException { if (comp instanceof JTable) { - StringSelection stringSelection = new StringSelection(getSelectedDataAtTable((JTable) comp)); + 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); @@ -191,13 +191,11 @@ public String getSelectedDataAtTable(JTable table) { selectData.append(table.getValueAt(row, 1).toString()).append("\n"); } - // 便于单行复制，去除最后一个换行符 if (!selectData.isEmpty()) { selectData.deleteCharAt(selectData.length() - 1); - return selectData.toString(); - } else { - return ""; } + + return selectData.toString(); } public JTable getDataTable() { diff --git a/src/main/java/hae/component/board/message/MessageRenderer.java b/src/main/java/hae/component/board/message/MessageRenderer.java index a613b0f..186efe2 100644 --- a/src/main/java/hae/component/board/message/MessageRenderer.java +++ b/src/main/java/hae/component/board/message/MessageRenderer.java @@ -4,16 +4,16 @@ import javax.swing.table.DefaultTableCellRenderer; import java.awt.*; import java.util.HashMap; -import java.util.List; +import java.util.LinkedList; import java.util.Map; public class MessageRenderer extends DefaultTableCellRenderer { - private final List log; + private final LinkedList log; private final Map colorMap = new HashMap<>(); private final JTable table; // 保存对表格的引用 - public MessageRenderer(List log, JTable table) { + public MessageRenderer(LinkedList log, JTable table) { this.log = log; // 与BurpSuite的颜色保持一致 this.colorMap.put("red", new Color(0xFF, 0x64, 0x64)); diff --git a/src/main/java/hae/component/board/message/MessageTableModel.java b/src/main/java/hae/component/board/message/MessageTableModel.java index cb94ec6..e0b9bb2 100644 --- a/src/main/java/hae/component/board/message/MessageTableModel.java +++ b/src/main/java/hae/component/board/message/MessageTableModel.java @@ -23,10 +23,6 @@ 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; @@ -35,16 +31,16 @@ public class MessageTableModel extends AbstractTableModel { private final MontoyaApi api; private final MessageTable messageTable; - private final JTabbedPane messageTab; private final JSplitPane splitPane; private final LinkedList log = new LinkedList<>(); private final LinkedList filteredLog; + private SwingWorker currentWorker; public MessageTableModel(MontoyaApi api) { this.filteredLog = new LinkedList<>(); this.api = api; - messageTab = new JTabbedPane(); + JTabbedPane messageTab = new JTabbedPane(); UserInterface userInterface = api.userInterface(); HttpRequestEditor requestViewer = userInterface.createHttpRequestEditor(READ_ONLY); HttpResponseEditor responseViewer = userInterface.createHttpResponseEditor(READ_ONLY); @@ -144,7 +140,11 @@ public void deleteByHost(String filterText) { filteredLog.clear(); List rowsToRemove = new ArrayList<>(); - new SwingWorker() { + if (currentWorker != null && !currentWorker.isDone()) { + currentWorker.cancel(true); + } + + currentWorker = new SwingWorker() { @Override protected Void doInBackground() { for (int i = 0; i < log.size(); i++) { @@ -164,53 +164,25 @@ protected Void doInBackground() { return null; } + }; - @Override - protected void done() { - if (!rowsToRemove.isEmpty()) { - int[] rows = rowsToRemove.stream().mapToInt(Integer::intValue).toArray(); - fireTableRowsDeleted(rows[0], rows[rows.length - 1]); - } - } - }.execute(); + currentWorker.execute(); } public void applyHostFilter(String filterText) { filteredLog.clear(); - 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(); + log.forEach(entry -> { + MessageEntry finalEntry = getEntryByFile(entry); + String host = StringProcessor.getHostByUrl(finalEntry.getUrl()); + if (!host.isEmpty()) { + if (StringProcessor.matchesHostPattern(host, filterText) || filterText.contains("*")) { + filteredLog.add(finalEntry); } } + }); - } catch (Exception e) { - api.logging().logToError("applyHostFilter: " + e.getMessage()); - } finally { - executor.shutdown(); - fireTableDataChanged(); - } + fireTableDataChanged(); } private MessageEntry getEntryByFile(MessageEntry entry) { @@ -244,12 +216,12 @@ private MessageEntry getEntryByFile(MessageEntry entry) { public void applyMessageFilter(String tableName, String filterText) { filteredLog.clear(); for (MessageEntry entry : log) { - HttpRequestResponse requestResponse = entry.getRequestResponse(); // 标志变量，表示是否满足过滤条件 AtomicBoolean isMatched = new AtomicBoolean(false); - HttpRequest httpRequest = entry.getRequestResponse().request(); - HttpResponse httpResponse = entry.getRequestResponse().response(); + HttpRequestResponse requestResponse = entry.getRequestResponse(); + HttpRequest httpRequest = requestResponse.request(); + HttpResponse httpResponse = requestResponse.response(); String requestString = new String(httpRequest.toByteArray().getBytes(), StandardCharsets.UTF_8); String requestBody = new String(httpRequest.body().getBytes(), StandardCharsets.UTF_8); @@ -263,7 +235,6 @@ public void applyMessageFilter(String tableName, String filterText) { .map(HttpHeader::toString) .collect(Collectors.joining("\n")); - MessageEntry finalEntry = entry; Config.globalRules.keySet().forEach(i -> { for (Object[] objects : Config.globalRules.get(i)) { String name = objects[1].toString(); @@ -271,7 +242,7 @@ public void applyMessageFilter(String tableName, String filterText) { String scope = objects[6].toString(); // 从注释中查看是否包含当前规则名，包含的再进行查询，有效减少无意义的检索时间 - if (finalEntry.getComment().contains(name)) { + if (entry.getComment().contains(name)) { if (name.equals(tableName)) { // 标志变量，表示当前规则是否匹配 boolean isMatch = false; @@ -410,7 +381,7 @@ public MessageTable getMessageTable() { return messageTable; } - public List getLogs() { + public LinkedList getLogs() { return log; } @@ -426,21 +397,27 @@ public int getColumnCount() { @Override public Object getValueAt(int rowIndex, int columnIndex) { - if (filteredLog.isEmpty()) { - return ""; + if (!filteredLog.isEmpty()) { + try { + MessageEntry messageEntry = filteredLog.get(rowIndex); + + if (messageEntry != null) { + return switch (columnIndex) { + case 0 -> messageEntry.getMethod(); + case 1 -> messageEntry.getUrl(); + case 2 -> messageEntry.getComment(); + case 3 -> messageEntry.getStatus(); + case 4 -> messageEntry.getLength(); + case 5 -> messageEntry.getColor(); + default -> ""; + }; + } + } catch (Exception e) { + api.logging().logToError("getValueAt: " + e.getMessage()); + } } - MessageEntry messageEntry = filteredLog.get(rowIndex); - - return switch (columnIndex) { - case 0 -> messageEntry.getMethod(); - case 1 -> messageEntry.getUrl(); - case 2 -> messageEntry.getComment(); - case 3 -> messageEntry.getStatus(); - case 4 -> messageEntry.getLength(); - case 5 -> messageEntry.getColor(); - default -> ""; - }; + return ""; } @Override @@ -459,8 +436,6 @@ public String getColumnName(int columnIndex) { public class MessageTable extends JTable { private MessageEntry messageEntry; private SwingWorker currentWorker; - // 设置响应报文返回的最大长度 - private final int MAX_LENGTH = 5242880; private int lastSelectedIndex = -1; private final HttpRequestEditor requestEditor; private final HttpResponseEditor responseEditor; @@ -477,7 +452,7 @@ public void changeSelection(int row, int col, boolean toggle, boolean extend) { requestEditor.setRequest(HttpRequest.httpRequest("Loading...")); responseEditor.setResponse(HttpResponse.httpResponse("Loading...")); - + if (currentWorker != null && !currentWorker.isDone()) { currentWorker.cancel(true); } @@ -495,11 +470,6 @@ protected Void doInBackground() { ByteArray requestByte = httpRequestResponse.request().toByteArray(); ByteArray responseByte = httpRequestResponse.response().toByteArray(); - if (responseByte.length() > MAX_LENGTH) { - String ellipsis = "\r\n......"; - responseByte = responseByte.subArray(0, MAX_LENGTH).withAppended(ellipsis); - } - requestEditor.setRequest(HttpRequest.httpRequest(messageEntry.getRequestResponse().httpService(), requestByte)); responseEditor.setResponse(HttpResponse.httpResponse(responseByte)); diff --git a/src/main/java/hae/instances/http/HttpMessageHandler.java b/src/main/java/hae/instances/http/HttpMessageHandler.java index f5f63fe..5f2d4e7 100644 --- a/src/main/java/hae/instances/http/HttpMessageHandler.java +++ b/src/main/java/hae/instances/http/HttpMessageHandler.java @@ -12,6 +12,7 @@ import hae.utils.ConfigLoader; import hae.utils.string.StringProcessor; +import javax.swing.*; import java.util.ArrayList; import java.util.Arrays; import java.util.List; @@ -45,19 +46,23 @@ public RequestToBeSentAction handleHttpRequestToBeSent(HttpRequestToBeSent httpR Annotations annotations = httpRequestToBeSent.annotations(); - httpRequest.set(httpRequestToBeSent); + try { + httpRequest.set(httpRequestToBeSent); - host.set(StringProcessor.getHostByUrl(httpRequestToBeSent.url())); + host.set(StringProcessor.getHostByUrl(httpRequestToBeSent.url())); - String[] hostList = configLoader.getBlockHost().split("\\|"); - boolean isBlockHost = RequestEditor.isBlockHost(hostList, host.get()); + String[] hostList = configLoader.getBlockHost().split("\\|"); + boolean isBlockHost = RequestEditor.isBlockHost(hostList, host.get()); - List suffixList = Arrays.asList(configLoader.getExcludeSuffix().split("\\|")); - matches.set(suffixList.contains(httpRequestToBeSent.fileExtension().toLowerCase()) || isBlockHost); + List suffixList = Arrays.asList(configLoader.getExcludeSuffix().split("\\|")); + matches.set(suffixList.contains(httpRequestToBeSent.fileExtension().toLowerCase()) || isBlockHost); - if (!matches.get()) { - List> result = messageProcessor.processRequest(host.get(), httpRequestToBeSent, true); - setColorAndCommentList(result); + if (!matches.get()) { + List> result = messageProcessor.processRequest(host.get(), httpRequestToBeSent, true); + setColorAndCommentList(result); + } + } catch (Exception e) { + api.logging().logToError("handleHttpRequestToBeSent: " + e.getMessage()); } return RequestToBeSentAction.continueWith(httpRequestToBeSent, annotations); @@ -85,7 +90,14 @@ public ResponseReceivedAction handleHttpResponseReceived(HttpResponseReceived ht String status = String.valueOf(httpResponseReceived.statusCode()); String length = String.valueOf(httpResponseReceived.toByteArray().length()); - messageTableModel.add(httpRequestResponse, url, method, status, length, comment, color, "", ""); + // 后台提交，防止线程阻塞 + new SwingWorker() { + @Override + protected Void doInBackground() { + messageTableModel.add(httpRequestResponse, url, method, status, length, comment, color, "", ""); + return null; + } + }.run(); } } @@ -93,7 +105,7 @@ public ResponseReceivedAction handleHttpResponseReceived(HttpResponseReceived ht } private void setColorAndCommentList(List> result) { - if (result != null && !result.isEmpty() && result.size() > 0) { + if (result != null && !result.isEmpty()) { colorList.get().add(result.get(0).get("color")); commentList.get().add(result.get(1).get("comment")); } diff --git a/src/main/java/hae/instances/http/utils/MessageProcessor.java b/src/main/java/hae/instances/http/utils/MessageProcessor.java index 84e7578..84b7524 100644 --- a/src/main/java/hae/instances/http/utils/MessageProcessor.java +++ b/src/main/java/hae/instances/http/utils/MessageProcessor.java @@ -71,7 +71,7 @@ private List> getDataList(Map> o List> highlightList = new ArrayList<>(); List> extractList = new ArrayList<>(); - if (obj != null && !obj.isEmpty() && obj.size() > 0) { + if (obj != null && !obj.isEmpty()) { if (actionFlag) { List> resultList = extractColorsAndComments(obj); List colorList = resultList.get(0); diff --git a/src/main/java/hae/utils/project/ProjectProcessor.java b/src/main/java/hae/utils/project/ProjectProcessor.java index 1f8efd1..a3371b4 100644 --- a/src/main/java/hae/utils/project/ProjectProcessor.java +++ b/src/main/java/hae/utils/project/ProjectProcessor.java @@ -11,10 +11,9 @@ import java.nio.file.Path; import java.nio.file.StandardCopyOption; import java.util.*; -import java.util.concurrent.ExecutionException; +import java.util.concurrent.Callable; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.Future; import java.util.zip.ZipEntry; import java.util.zip.ZipFile; import java.util.zip.ZipOutputStream; @@ -27,8 +26,9 @@ public ProjectProcessor(MontoyaApi api) { } public boolean createHaeFile(String haeFilePath, String host, Map> dataMap, Map> urlMap, Map> httpMap) { - ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors()); - List> futures = new ArrayList<>(); + ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2); + + List> tasks = new ArrayList<>(); ByteArrayOutputStream dataYamlStream = new ByteArrayOutputStream(); ByteArrayOutputStream urlYamlStream = new ByteArrayOutputStream(); @@ -52,7 +52,7 @@ public boolean createHaeFile(String haeFilePath, String host, Map httpItem = httpMap.get(httpHash); - futures.add(executorService.submit(() -> { + tasks.add(() -> { try { ByteArrayOutputStream httpOutStream = new ByteArrayOutputStream(); byte[] request = (byte[]) httpItem.get("request"); @@ -69,29 +69,25 @@ public boolean createHaeFile(String haeFilePath, String host, Map future : futures) { - try { - future.get(); - } catch (InterruptedException | ExecutionException e) { - Thread.currentThread().interrupt(); - } + return null; + }); } + + executor.invokeAll(tasks); } catch (Exception e) { api.logging().logToError("createHaeFile: " + e.getMessage()); return false; } finally { - executorService.shutdown(); + executor.shutdown(); } return true; } public HaeFileContent readHaeFile(String haeFilePath) { - ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors()); - List> futures = new ArrayList<>(); + ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2); + List> tasks = new ArrayList<>(); HaeFileContent haeFileContent = new HaeFileContent(api); LoaderOptions loaderOptions = new LoaderOptions(); @@ -112,13 +108,16 @@ public HaeFileContent readHaeFile(String haeFilePath) { String fileName = entry.getName(); if (fileName.startsWith("http/")) { Path filePath = tempDirectory.resolve(fileName.substring("http/".length())); - futures.add(executorService.submit(() -> { + + tasks.add(() -> { try (InputStream in = zipFile.getInputStream(entry)) { Files.copy(in, filePath, StandardCopyOption.REPLACE_EXISTING); } catch (IOException e) { api.logging().logToError("readHaeFile: " + e.getMessage()); } - })); + + return null; + }); } else { try (InputStream in = zipFile.getInputStream(entry)) { switch (fileName) { @@ -133,13 +132,7 @@ public HaeFileContent readHaeFile(String haeFilePath) { } } - for (Future future : futures) { - try { - future.get(); - } catch (InterruptedException | ExecutionException e) { - Thread.currentThread().interrupt(); - } - } + executor.invokeAll(tasks); } } } catch (Exception e) { @@ -149,7 +142,7 @@ public HaeFileContent readHaeFile(String haeFilePath) { } haeFileContent = null; } finally { - executorService.shutdown(); + executor.shutdown(); } return haeFileContent; diff --git a/src/main/java/hae/utils/string/StringProcessor.java b/src/main/java/hae/utils/string/StringProcessor.java index 7168b75..6588f41 100644 --- a/src/main/java/hae/utils/string/StringProcessor.java +++ b/src/main/java/hae/utils/string/StringProcessor.java @@ -119,17 +119,6 @@ public static String getHostByUrl(String url) { 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"); }