Version: 4.0 Update

2024-12-21 15:34:45 +08:00
parent 1f1ca99f10
commit daacb2e146
24 changed files with 469 additions and 1427 deletions
--- a/src/main/java/hae/utils/DataManager.java
+++ b/src/main/java/hae/utils/DataManager.java
@@ -0,0 +1,85 @@
+package hae.utils;
+
+import burp.api.montoya.MontoyaApi;
+import burp.api.montoya.http.message.HttpRequestResponse;
+import burp.api.montoya.http.message.requests.HttpRequest;
+import burp.api.montoya.http.message.responses.HttpResponse;
+import burp.api.montoya.persistence.PersistedList;
+import burp.api.montoya.persistence.PersistedObject;
+import burp.api.montoya.persistence.Persistence;
+import hae.component.board.message.MessageTableModel;
+import hae.instances.http.utils.RegularMatcher;
+
+public class DataManager {
+    private final MontoyaApi api;
+    private final Persistence persistence;
+
+    public DataManager(MontoyaApi api) {
+        this.api = api;
+        this.persistence = api.persistence();
+    }
+
+    private void saveIndex(String indexName, String indexValue) {
+        PersistedList<String> indexList = persistence.extensionData().getStringList(indexName);
+
+        if (indexList != null && !indexList.isEmpty()) {
+            persistence.extensionData().deleteStringList(indexName);
+        } else {
+            indexList = PersistedList.persistedStringList();
+        }
+
+        if (!indexList.contains(indexValue)) {
+            indexList.add(indexValue);
+        }
+
+        persistence.extensionData().setStringList(indexName, indexList);
+    }
+
+    public void putData(String dataType, String dataName, PersistedObject persistedObject) {
+        if (persistence.extensionData().getChildObject(dataName) != null) {
+            persistence.extensionData().deleteChildObject(dataName);
+        }
+        persistence.extensionData().setChildObject(dataName, persistedObject);
+        saveIndex(dataType, dataName);
+    }
+
+    public void loadData(MessageTableModel messageTableModel) {
+        // 1. 获取索引
+        PersistedList<String> dataIndex = persistence.extensionData().getStringList("data"); // 数据索引
+        PersistedList<String> messageIndex = persistence.extensionData().getStringList("message"); // 消息索引
+
+        // 2. 从索引获取数据
+        loadHaEData(dataIndex);
+        loadMessageData(messageIndex, messageTableModel);
+
+    }
+
+    private void loadHaEData(PersistedList<String> dataIndex) {
+        if (dataIndex != null && !dataIndex.isEmpty()) {
+            dataIndex.parallelStream().forEach(index -> {
+                PersistedObject dataObj = persistence.extensionData().getChildObject(index);
+                dataObj.stringListKeys().forEach(dataKey -> {
+                    RegularMatcher.putDataToGlobalMap(api, index, dataKey, dataObj.getStringList(dataKey).stream().toList(), false);
+                });
+            });
+        }
+    }
+
+    private void loadMessageData(PersistedList<String> messageIndex, MessageTableModel messageTableModel) {
+        if (messageIndex != null && !messageIndex.isEmpty()) {
+            messageIndex.parallelStream().forEach(index -> {
+                PersistedObject dataObj = persistence.extensionData().getChildObject(index);
+                HttpRequestResponse messageInfo = dataObj.getHttpRequestResponse("messageInfo");
+                String comment = dataObj.getString("comment");
+                String color = dataObj.getString("color");
+                HttpRequest request = messageInfo.request();
+                HttpResponse response = messageInfo.response();
+                String method = request.method();
+                String url = request.url();
+                String status = String.valueOf(response.statusCode());
+                String length = String.valueOf(response.toByteArray().length());
+                messageTableModel.add(messageInfo, url, method, status, length, comment, color, false);
+            });
+        }
+    }
+}