Version: 3.1 Update

src/main/java/hae/utils/config/ConfigLoader.java

@@ -1,6 +1,7 @@
 package hae.utils.config;
 
 import burp.api.montoya.MontoyaApi;
+import burp.api.montoya.http.RequestOptions;
 import burp.api.montoya.http.message.HttpRequestResponse;
 import burp.api.montoya.http.message.requests.HttpRequest;
 import hae.Config;
@@ -11,6 +12,7 @@ import org.yaml.snakeyaml.representer.Representer;
 import java.io.*;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
+import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.*;
 
@@ -44,7 +46,7 @@ public class ConfigLoader {
 
         File rulesFilePath = new File(this.rulesFilePath);
         if (!(rulesFilePath.exists() && rulesFilePath.isFile())) {
-            initRules();
+            initRulesByRes();
         }
 
         Config.globalRules = getRules();
@@ -76,6 +78,7 @@ public class ConfigLoader {
     public void initConfig() {
         Map<String, Object> r = new LinkedHashMap<>();
         r.put("excludeSuffix", getExcludeSuffix());
+        r.put("blockHost", getBlockHost());
         try {
             Writer ws = new OutputStreamWriter(Files.newOutputStream(Paths.get(configFilePath)), StandardCharsets.UTF_8);
             yaml.dump(r, ws);
@@ -88,24 +91,6 @@ public class ConfigLoader {
         return rulesFilePath;
     }
 
-    public String getExcludeSuffix() {
-        File yamlSetting = new File(configFilePath);
-        if (!yamlSetting.exists() || !yamlSetting.isFile()) {
-            return Config.suffix;
-        }
-
-        try (InputStream inorder = Files.newInputStream(Paths.get(configFilePath))) {
-            Map<String, Object> r = new Yaml().load(inorder);
-
-            if (r.containsKey("excludeSuffix")) {
-                return r.get("excludeSuffix").toString();
-            }
-        } catch (Exception ignored) {
-        }
-
-        return Config.suffix;
-    }
-
     // 获取规则配置
     public Map<String, Object[][]> getRules() {
         Map<String, Object[][]> rules = new HashMap<>();
@@ -149,18 +134,104 @@ public class ConfigLoader {
         return rules;
     }
 
-    public void setExcludeSuffix(String excludeSuffix) {
-        Map<String, Object> r = new LinkedHashMap<>();
-        r.put("excludeSuffix", excludeSuffix);
-        try {
-            Writer ws = new OutputStreamWriter(Files.newOutputStream(Paths.get(configFilePath)), StandardCharsets.UTF_8);
-            yaml.dump(r, ws);
-            ws.close();
+    public String getBlockHost() {
+        File yamlSetting = new File(configFilePath);
+        if (!yamlSetting.exists() || !yamlSetting.isFile()) {
+            return Config.host;
+        }
+
+        try (InputStream inorder = Files.newInputStream(Paths.get(configFilePath))) {
+            Map<String, Object> r = new Yaml().load(inorder);
+
+            if (r.containsKey("blockHost")) {
+                return r.get("blockHost").toString();
+            }
         } catch (Exception ignored) {
         }
+
+        return Config.host;
+    }
+
+    public String getExcludeSuffix() {
+        File yamlSetting = new File(configFilePath);
+        if (!yamlSetting.exists() || !yamlSetting.isFile()) {
+            return Config.suffix;
+        }
+
+        try (InputStream inorder = Files.newInputStream(Paths.get(configFilePath))) {
+            Map<String, Object> r = new Yaml().load(inorder);
+
+            if (r.containsKey("excludeSuffix")) {
+                return r.get("excludeSuffix").toString();
+            }
+        } catch (Exception ignored) {
+        }
+
+        return Config.suffix;
+    }
+
+    private Map<String, Object> loadCurrentConfig() {
+        Path path = Paths.get(configFilePath);
+        if (!Files.exists(path)) {
+            return new LinkedHashMap<>(); // 返回空的Map，表示没有当前配置
+        }
+
+        try (InputStream in = Files.newInputStream(path)) {
+            return yaml.load(in);
+        } catch (IOException e) {
+            return new LinkedHashMap<>(); // 读取失败时也返回空的Map
+        }
     }
 
-    public void initRules() {
+    public void setExcludeSuffix(String excludeSuffix) {
+        Map<String, Object> currentConfig = loadCurrentConfig();
+        currentConfig.put("excludeSuffix", excludeSuffix); // 更新配置
+
+        try (Writer ws = new OutputStreamWriter(Files.newOutputStream(Paths.get(configFilePath)), StandardCharsets.UTF_8)) {
+            yaml.dump(currentConfig, ws);
+        } catch (IOException ignored) {
+        }
+    }
+
+    public void setBlockHost(String blockHost) {
+        Map<String, Object> currentConfig = loadCurrentConfig();
+        currentConfig.put("blockHost", blockHost); // 更新配置
+
+        try (Writer ws = new OutputStreamWriter(Files.newOutputStream(Paths.get(configFilePath)), StandardCharsets.UTF_8)) {
+            yaml.dump(currentConfig, ws);
+        } catch (IOException ignored) {
+        }
+    }
+
+    public void initRulesByRes() {
+        boolean isCopySuccess = copyRulesToFile(this.rulesFilePath);
+        if (!isCopySuccess) {
+            api.extension().unload();
+        }
+    }
+
+    private boolean copyRulesToFile(String targetFilePath) {
+        InputStream inputStream = getClass().getClassLoader().getResourceAsStream("rules/Rules.yml");
+        File targetFile = new File(targetFilePath);
+
+        try (inputStream; OutputStream outputStream = new FileOutputStream(targetFile)) {
+            if (inputStream != null) {
+                byte[] buffer = new byte[1024];
+                int length;
+
+                while ((length = inputStream.read(buffer)) > 0) {
+                    outputStream.write(buffer, 0, length);
+                }
+
+                return true;
+            }
+        } catch (IOException ignored) {
+        }
+
+        return false;
+    }
+
+    public void initRulesByNet() {
         Thread t = new Thread() {
             public void run() {
                 pullRules();
@@ -177,7 +248,7 @@ public class ConfigLoader {
         try {
             String url = "https://raw.githubusercontent.com/gh0stkey/HaE/gh-pages/Rules.yml";
             HttpRequest httpRequest = HttpRequest.httpRequestFromUrl(url);
-            HttpRequestResponse requestResponse = api.http().sendRequest(httpRequest);
+            HttpRequestResponse requestResponse = api.http().sendRequest(httpRequest, RequestOptions.requestOptions().withUpstreamTLSVerification());
             String responseBody = requestResponse.response().bodyToString();
             if (responseBody.contains("rules")) {
                 FileOutputStream fileOutputStream = new FileOutputStream(rulesFilePath);