admin/HaE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Version: 3.2.2 Update

Browse Source
This commit is contained in:
gh0stkey
2024-06-19 22:16:57 +08:00
parent 54973d9f4f
commit e5f55b6c4c
9 changed files with 190 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/main/java/hae/HaE.java
View File

@@ -2,7 +2,9 @@ package hae;
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 class HaE implements BurpExtension {
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();
}
});
}
}

147
src/main/java/hae/component/board/Databoard.java
View File

@@ -23,7 +23,10 @@ import java.awt.event.*;
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<Boolean, Void> currentWorker;
private SwingWorker<Boolean, Void> handleComboBoxWorker;
private SwingWorker<Void, Void> applyHostFilterWorker;
private SwingWorker<List<String>, Void> exportActionWorker;
private SwingWorker<List<String>, Void> importActionWorker;
public Databoard(MontoyaApi api, ConfigLoader configLoader, MessageTableModel messageTableModel) {
this.api = api;
@@ -77,11 +83,10 @@ public class Databoard extends JPanel {
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 @@ public class Databoard extends JPanel {
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<Boolean, Void>() {
@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();
handleComboBoxWorker = new SwingWorker<Boolean, Void>() {
@Override
protected Boolean doInBackground() {
return populateTabbedPaneByHost(selectedHost);
}
splitPane.setVisible(true);
hostTextField.setText(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();
hostComboBox.setPopupVisible(false);
applyHostFilter(selectedHost);
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 @@ public class Databoard extends JPanel {
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<String, List<String>> ruleMap = dataMap.get(key);
for (String ruleKey : ruleMap.keySet()) {
@@ -286,6 +295,7 @@ public class Databoard extends JPanel {
return true;
}
return false;
}
@@ -316,7 +326,11 @@ public class Databoard extends JPanel {
TableRowSorter<TableModel> sorter = (TableRowSorter<TableModel>) messageTable.getRowSorter();
String cleanedText = StringProcessor.replaceFirstOccurrence(filterText, "*.", "");
new SwingWorker<Void, Void>() {
if (applyHostFilterWorker != null && !applyHostFilterWorker.isDone()) {
applyHostFilterWorker.cancel(true);
}
applyHostFilterWorker = new SwingWorker<Void, Void>() {
@Override
protected Void doInBackground() throws Exception {
RowFilter<Object, Object> rowFilter = new RowFilter<Object, Object>() {
@@ -340,8 +354,9 @@ public class Databoard extends JPanel {
protected void done() {
setProgressBar(false);
}
}.execute();
};
applyHostFilterWorker.execute();
}
private List<String> getHostByList() {
@@ -364,7 +379,11 @@ public class Databoard extends JPanel {
return;
}
new SwingWorker<List<String>, Void>() {
if (exportActionWorker != null && !exportActionWorker.isDone()) {
exportActionWorker.cancel(true);
}
exportActionWorker = new SwingWorker<List<String>, Void>() {
@Override
protected List<String> doInBackground() {
ConcurrentHashMap<String, Map<String, List<String>>> dataMap = Config.globalDataMap;
@@ -382,7 +401,9 @@ public class Databoard extends JPanel {
} catch (Exception ignored) {
}
}
}.execute();
};
exportActionWorker.execute();
}
private List<String> exportData(String selectedHost, String exportDir, Map<String, Map<String, List<String>>> dataMap) {
@@ -469,7 +490,11 @@ public class Databoard extends JPanel {
return;
}
new SwingWorker<List<String>, Void>() {
if (importActionWorker != null && !importActionWorker.isDone()) {
importActionWorker.cancel(true);
}
importActionWorker = new SwingWorker<List<String>, Void>() {
@Override
protected List<String> doInBackground() {
List<String> filesWithExtension = findFilesWithExtension(new File(exportDir), ".hae");
@@ -489,15 +514,18 @@ public class Databoard extends JPanel {
} 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<Future<?>> futures = new ArrayList<>();
List<Callable<Void>> tasks = new ArrayList<>();
if (readStatus) {
try {
@@ -505,7 +533,7 @@ public class Databoard extends JPanel {
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 @@ public class Databoard extends JPanel {
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 @@ public class Databoard extends JPanel {
}
});
if (!StringProcessor.matchHostIsIp(host) && !host.contains("*.")) {
String baseDomain = StringProcessor.getBaseDomain(StringProcessor.extractHostname(host));
// 删除无用的数据
Set<String> 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<String> 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<String> keysToRemove = wildcardKeys.stream()
.filter(key -> !existingSuffixes.contains(key.substring(1)))
.collect(Collectors.toSet());
if (Config.globalDataMap.keySet().size() == 1 && Config.globalDataMap.keySet().stream().anyMatch(key -> key.contains("*"))) {
keysToRemove.forEach(Config.globalDataMap::remove);
if (Config.globalDataMap.keySet().size() == 1 && Config.globalDataMap.keySet().stream().anyMatch(key -> key.equals("*"))) {
Config.globalDataMap.keySet().remove("*");
}

8
src/main/java/hae/component/board/Datatable.java
View File

@@ -158,7 +158,7 @@ public class Datatable extends JPanel {
@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 class Datatable extends JPanel {
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() {

6
src/main/java/hae/component/board/message/MessageRenderer.java
View File

@@ -4,16 +4,16 @@ import javax.swing.*;
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<MessageEntry> log;
private final LinkedList<MessageEntry> log;
private final Map<String, Color> colorMap = new HashMap<>();
private final JTable table; // 保存对表格的引用
public MessageRenderer(List<MessageEntry> log, JTable table) {
public MessageRenderer(LinkedList<MessageEntry> log, JTable table) {
this.log = log;
// 与BurpSuite的颜色保持一致
this.colorMap.put("red", new Color(0xFF, 0x64, 0x64));

114
src/main/java/hae/component/board/message/MessageTableModel.java
View File

@@ -23,10 +23,6 @@ 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;
@@ -35,16 +31,16 @@ import static burp.api.montoya.ui.editor.EditorOptions.READ_ONLY;
public class MessageTableModel extends AbstractTableModel {
private final MontoyaApi api;
private final MessageTable messageTable;
private final JTabbedPane messageTab;
private final JSplitPane splitPane;
private final LinkedList<MessageEntry> log = new LinkedList<>();
private final LinkedList<MessageEntry> filteredLog;
private SwingWorker<Void, Void> 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 class MessageTableModel extends AbstractTableModel {
filteredLog.clear();
List<Integer> rowsToRemove = new ArrayList<>();
new SwingWorker<Void, Void>() {
if (currentWorker != null && !currentWorker.isDone()) {
currentWorker.cancel(true);
}
currentWorker = new SwingWorker<Void, Void>() {
@Override
protected Void doInBackground() {
for (int i = 0; i < log.size(); i++) {
@@ -164,53 +164,25 @@ public class MessageTableModel extends AbstractTableModel {
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<Future<?>> 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 @@ public class MessageTableModel extends AbstractTableModel {
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 class MessageTableModel extends AbstractTableModel {
.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 class MessageTableModel extends AbstractTableModel {
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 class MessageTableModel extends AbstractTableModel {
return messageTable;
}
public List<MessageEntry> getLogs() {
public LinkedList<MessageEntry> getLogs() {
return log;
}
@@ -426,21 +397,27 @@ public class MessageTableModel extends AbstractTableModel {
@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 class MessageTableModel extends AbstractTableModel {
public class MessageTable extends JTable {
private MessageEntry messageEntry;
private SwingWorker<Object, Void> currentWorker;
// 设置响应报文返回的最大长度
private final int MAX_LENGTH = 5242880;
private int lastSelectedIndex = -1;
private final HttpRequestEditor requestEditor;
private final HttpResponseEditor responseEditor;
@@ -477,7 +452,7 @@ public class MessageTableModel extends AbstractTableModel {
requestEditor.setRequest(HttpRequest.httpRequest("Loading..."));
responseEditor.setResponse(HttpResponse.httpResponse("Loading..."));
if (currentWorker != null && !currentWorker.isDone()) {
currentWorker.cancel(true);
}
@@ -495,11 +470,6 @@ public class MessageTableModel extends AbstractTableModel {
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));

34
src/main/java/hae/instances/http/HttpMessageHandler.java
View File

@@ -12,6 +12,7 @@ import hae.instances.http.utils.MessageProcessor;
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 class HttpMessageHandler implements HttpHandler {
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<String> suffixList = Arrays.asList(configLoader.getExcludeSuffix().split("\\|"));
matches.set(suffixList.contains(httpRequestToBeSent.fileExtension().toLowerCase()) || isBlockHost);
List<String> suffixList = Arrays.asList(configLoader.getExcludeSuffix().split("\\|"));
matches.set(suffixList.contains(httpRequestToBeSent.fileExtension().toLowerCase()) || isBlockHost);
if (!matches.get()) {
List<Map<String, String>> result = messageProcessor.processRequest(host.get(), httpRequestToBeSent, true);
setColorAndCommentList(result);
if (!matches.get()) {
List<Map<String, String>> 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 class HttpMessageHandler implements HttpHandler {
String status = String.valueOf(httpResponseReceived.statusCode());
String length = String.valueOf(httpResponseReceived.toByteArray().length());
messageTableModel.add(httpRequestResponse, url, method, status, length, comment, color, "", "");
// 后台提交，防止线程阻塞
new SwingWorker<Void, Void>() {
@Override
protected Void doInBackground() {
messageTableModel.add(httpRequestResponse, url, method, status, length, comment, color, "", "");
return null;
}
}.run();
}
}
@@ -93,7 +105,7 @@ public class HttpMessageHandler implements HttpHandler {
}
private void setColorAndCommentList(List<Map<String, String>> 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"));
}

2
src/main/java/hae/instances/http/utils/MessageProcessor.java
View File

@@ -71,7 +71,7 @@ public class MessageProcessor {
List<Map<String, String>> highlightList = new ArrayList<>();
List<Map<String, String>> extractList = new ArrayList<>();
if (obj != null && !obj.isEmpty() && obj.size() > 0) {
if (obj != null && !obj.isEmpty()) {
if (actionFlag) {
List<List<String>> resultList = extractColorsAndComments(obj);
List<String> colorList = resultList.get(0);

45
src/main/java/hae/utils/project/ProjectProcessor.java
View File

@@ -11,10 +11,9 @@ 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.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 class ProjectProcessor {
}
public boolean createHaeFile(String haeFilePath, String host, Map<String, List<String>> dataMap, Map<String, Map<String, Object>> urlMap, Map<String, Map<String, Object>> httpMap) {
ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
List<Future<?>> futures = new ArrayList<>();
ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2);
List<Callable<Void>> tasks = new ArrayList<>();
ByteArrayOutputStream dataYamlStream = new ByteArrayOutputStream();
ByteArrayOutputStream urlYamlStream = new ByteArrayOutputStream();
@@ -52,7 +52,7 @@ public class ProjectProcessor {
for (String httpHash : httpMap.keySet()) {
Map<String, Object> 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 class ProjectProcessor {
} catch (Exception e) {
api.logging().logToError("createHaeFile: " + e.getMessage());
}
}));
return null;
});
}
for (Future<?> future : futures) {
try {
future.get();
} catch (InterruptedException | ExecutionException e) {
Thread.currentThread().interrupt();
}
}
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<Future<?>> futures = new ArrayList<>();
ExecutorService executor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors() * 2);
List<Callable<Void>> tasks = new ArrayList<>();
HaeFileContent haeFileContent = new HaeFileContent(api);
LoaderOptions loaderOptions = new LoaderOptions();
@@ -112,13 +108,16 @@ public class ProjectProcessor {
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 class ProjectProcessor {
}
}
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 class ProjectProcessor {
}
haeFileContent = null;
} finally {
executorService.shutdown();
executor.shutdown();
}
return haeFileContent;

11
src/main/java/hae/utils/string/StringProcessor.java
View File

@@ -119,17 +119,6 @@ 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");
}