Version: 2.0.7 Update

gh0stkey · Sep 7, 2021 · 5c326d3 · 5c326d3
1 parent 62edae0
commit 5c326d3
Show file tree

Hide file tree

Showing 8 changed files with 53 additions and 75 deletions.
diff --git a/src/main/java/burp/BurpExtender.java b/src/main/java/burp/BurpExtender.java
@@ -6,6 +6,7 @@
 import javax.swing.*;
 import java.awt.*;
 import java.io.UnsupportedEncodingException;
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.List;
 import java.io.PrintWriter;
@@ -16,11 +17,10 @@
  */
 
 public class BurpExtender implements IBurpExtender, IHttpListener, IMessageEditorTabFactory, ITab {
-    private MainUI main = new MainUI();
+    private final MainUI main = new MainUI();
     private static PrintWriter stdout;
     private IBurpExtenderCallbacks callbacks;
     private static IExtensionHelpers helpers;
-    private static IMessageEditorTab HaETab;
     MatchHTTP mh = new MatchHTTP();
     ExtractContent ec = new ExtractContent();
     DoAction da = new DoAction();
@@ -33,20 +33,15 @@ public void registerExtenderCallbacks(final IBurpExtenderCallbacks callbacks)
         this.callbacks = callbacks;
         BurpExtender.helpers = callbacks.getHelpers();
 
-        String version = "2.0.6";
+        String version = "2.0.7";
         callbacks.setExtensionName(String.format("HaE (%s) - Highlighter and Extractor", version));
         // 定义输出
         stdout = new PrintWriter(callbacks.getStdout(), true);
         stdout.println("@UI Author: 0chencc");
         stdout.println("@Core Author: EvilChen");
         stdout.println("@Github: https://github.com/gh0stkey/HaE");
         // UI
-        SwingUtilities.invokeLater(new Runnable() {
-            @Override
-            public void run() {
-                initialize();
-            }
-        });
+        SwingUtilities.invokeLater(this::initialize);
 
         callbacks.registerHttpListener(BurpExtender.this);
         callbacks.registerMessageEditorTabFactory(BurpExtender.this);
@@ -109,15 +104,15 @@ public void processHttpMessage(int toolFlag, boolean messageIsRequest, IHttpRequ
 
             List<String> colorList = da.highlightList(obj);
             if (colorList.size() != 0) {
-                String color = uc.getEndColor(gck.getColorKeys(colorList, Config.colorArray), Config.colorArray);
+                String color = uc.getEndColor(gck.getColorKeys(colorList));
                 messageInfo.setHighlight(color);
             }
         }
 
     }
 
     class MarkInfoTab implements IMessageEditorTab {
-        private ITextEditor markInfoText;
+        private final ITextEditor markInfoText;
         private byte[] currentMessage;
         private final IMessageEditorController controller;
         private byte[] extractRequestContent;
@@ -209,11 +204,7 @@ public byte[] getSelectedData() {
          */
         @Override
         public void setMessage(byte[] content, boolean isRequest) {
-            try {
-                String c = new String(content, "UTF-8").intern();
-            } catch (UnsupportedEncodingException e) {
-                stdout.println(e);
-            }
+            String c = new String(content, StandardCharsets.UTF_8).intern();
             if (content.length > 0) {
                 if (isRequest) {
                     markInfoText.setText(extractRequestContent);
@@ -227,7 +218,6 @@ public void setMessage(byte[] content, boolean isRequest) {
 
     @Override
     public IMessageEditorTab createNewInstance(IMessageEditorController controller, boolean editable) {
-        HaETab = new MarkInfoTab(controller, editable);
-        return HaETab;
+        return new MarkInfoTab(controller, editable);
     }
 }
diff --git a/src/main/java/burp/action/DoAction.java b/src/main/java/burp/action/DoAction.java
@@ -22,7 +22,7 @@ public String extractString(Map<String, Map<String, Object>> obj) {
     }
 
     public List<String> highlightList(Map<String, Map<String, Object>> obj) {
-        List<String> colorList = new ArrayList<String>();
+        List<String> colorList = new ArrayList<>();
         obj.keySet().forEach(i->{
             Map<String, Object> tmpMap = obj.get(i);
             String color = tmpMap.get("color").toString();

diff --git a/src/main/java/burp/action/ExtractContent.java b/src/main/java/burp/action/ExtractContent.java
@@ -1,6 +1,6 @@
 package burp.action;
 
-import java.io.UnsupportedEncodingException;
+import java.nio.charset.StandardCharsets;
 import java.util.*;
 
 import dk.brics.automaton.Automaton;
@@ -18,12 +18,11 @@
  */
 
 public class ExtractContent {
-    private LoadConfigFile lcf = new LoadConfigFile();
-    private LoadRule lr = new LoadRule(lcf.getConfigPath());
 
     public Map<String, Map<String, Object>> matchRegex(byte[] content, String headers, byte[] body, String scopeString) {
         Map<String, Map<String, Object>> map = new HashMap<>(); // 最终返回的结果
-        Map<String,Object[][]> rules = lr.getConfig();
+        new LoadRule(LoadConfigFile.getConfigPath());
+        Map<String,Object[][]> rules = LoadRule.getConfig();
         rules.keySet().forEach(i -> {
             String matchContent = "";
             for (Object[] objects : rules.get(i)) {
@@ -43,23 +42,15 @@ public Map<String, Map<String, Object>> matchRegex(byte[] content, String header
                         case "any":
                         case "request":
                         case "response":
-                            try {
-                                matchContent = new String(content, "UTF-8").intern();
-                            } catch (UnsupportedEncodingException e) {
-                                e.printStackTrace();
-                            }
+                            matchContent = new String(content, StandardCharsets.UTF_8).intern();
                             break;
                         case "request header":
                         case "response header":
                             matchContent = headers;
                             break;
                         case "request body":
                         case "response body":
-                            try {
-                                matchContent = new String(body, "UTF-8").intern();
-                            } catch (UnsupportedEncodingException e) {
-                                e.printStackTrace();
-                            }
+                            matchContent = new String(body, StandardCharsets.UTF_8).intern();
                             break;
                     }
 
@@ -72,8 +63,8 @@ public Map<String, Map<String, Object>> matchRegex(byte[] content, String header
                             result.add(matcher.group(1));
                         }
                     } else {
-                        RegExp regexpr = new RegExp(regex);
-                        Automaton auto = regexpr.toAutomaton();
+                        RegExp regexp = new RegExp(regex);
+                        Automaton auto = regexp.toAutomaton();
                         RunAutomaton runAuto = new RunAutomaton(auto, true);
                         AutomatonMatcher autoMatcher = runAuto.newMatcher(matchContent);
                         while (autoMatcher.find()) {

diff --git a/src/main/java/burp/action/GetColorKey.java b/src/main/java/burp/action/GetColorKey.java
@@ -1,5 +1,6 @@
 package burp.action;
 
+import burp.Config;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -11,13 +12,14 @@ public class GetColorKey {
     /*
      * 颜色下标获取
      */
-    public List<Integer> getColorKeys(List<String> keys, String[] colorArray){
-        List<Integer> result = new ArrayList<Integer>();
+    public List<Integer> getColorKeys(List<String> keys){
+        List<Integer> result = new ArrayList<>();
+        String[] colorArray = Config.colorArray;
         int size = colorArray.length;
         // 根据颜色获取下标
-        for (int x = 0; x < keys.size(); x++) {
+        for (String key : keys) {
             for (int v = 0; v < size; v++) {
-                if (colorArray[v].equals(keys.get(x))) {
+                if (colorArray[v].equals(key)) {
                     result.add(v);
                 }
             }

diff --git a/src/main/java/burp/action/MatchHTTP.java b/src/main/java/burp/action/MatchHTTP.java
@@ -15,10 +15,6 @@ public class MatchHTTP {
     public boolean matchSuffix(String str) {
         Pattern pattern = new Pattern(String.format("[\\w]+[\\.](%s)", lc.getExcludeSuffix()), REFlags.IGNORE_CASE);
         Matcher matcher = pattern.matcher(str);
-        if(matcher.find()){
-            return true;
-        }else{
-            return false;
-        }
+        return matcher.find();
     }
 }
diff --git a/src/main/java/burp/action/UpgradeColor.java b/src/main/java/burp/action/UpgradeColor.java
@@ -1,9 +1,8 @@
 package burp.action;
 
-import java.util.ArrayList;
-import java.util.Comparator;
-import java.util.HashSet;
-import java.util.List;
+import burp.Config;
+
+import java.util.*;
 
 /*
  * @author EvilChen
@@ -14,43 +13,40 @@ public class UpgradeColor {
     /*
      * 颜色升级递归算法
      */
-    private String colorUpgrade(List<Integer> colorList, String[] colorArray) {
+    private void colorUpgrade(List<Integer> colorList) {
         int colorSize = colorList.size();
+        String[] colorArray = Config.colorArray;
         colorList.sort(Comparator.comparingInt(Integer::intValue));
         int i = 0;
-        List<Integer> stack = new ArrayList<Integer>();
+        List<Integer> stack = new ArrayList<>();
         while (i < colorSize) {
             if (stack.isEmpty()) {
                 stack.add(colorList.get(i));
-                i++;
             } else {
-                if (colorList.get(i) != stack.stream().reduce((first, second) -> second).orElse(99999999)) {
+                if (!Objects.equals(colorList.get(i), stack.stream().reduce((first, second) -> second).orElse(99999999))) {
                     stack.add(colorList.get(i));
-                    i++;
                 } else {
                     stack.set(stack.size() - 1, stack.get(stack.size() - 1) - 1);
-                    i++;
                 }
             }
-
+            i++;
         }
         // 利用HashSet删除重复元素
         HashSet tmpList = new HashSet(stack);
         if (stack.size() == tmpList.size()) {
             stack.sort(Comparator.comparingInt(Integer::intValue));
-            if(stack.get(0).equals(-1)) {
+            if(stack.get(0) < 0) {
                 this.endColor = colorArray[0];
             } else {
                 this.endColor = colorArray[stack.get(0)];
             }
         } else {
-            this.colorUpgrade(stack, colorArray);
+            this.colorUpgrade(stack);
         }
-        return "";
     }
 
-    public String getEndColor(List<Integer> colorList, String[] colorArray) {
-        colorUpgrade(colorList, colorArray);
+    public String getEndColor(List<Integer> colorList) {
+        colorUpgrade(colorList);
         return endColor;
     }
 }
diff --git a/src/main/java/burp/yaml/LoadConfigFile.java b/src/main/java/burp/yaml/LoadConfigFile.java
@@ -3,7 +3,10 @@
 import org.jetbrains.annotations.NotNull;
 import org.yaml.snakeyaml.Yaml;
 
+import burp.Config;
+
 import java.io.*;
+import java.nio.charset.StandardCharsets;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -12,7 +15,7 @@
  */
 
 public class LoadConfigFile {
-    private static Yaml yaml = new Yaml();
+    private static final Yaml yaml = new Yaml();
     private static final String SettingPath = "Setting.yml";
     private static final String ConfigPath = "Config.yml";
 
@@ -28,7 +31,7 @@ public void init(){
             r.put("configPath", ConfigPath);
             r.put("excludeSuffix", getExcludeSuffix());
             try{
-                Writer ws = new OutputStreamWriter(new FileOutputStream(SettingPath),"UTF-8");
+                Writer ws = new OutputStreamWriter(new FileOutputStream(SettingPath), StandardCharsets.UTF_8);
                 yaml.dump(r, ws);
             }catch (Exception ex){
                 ex.printStackTrace();
@@ -44,11 +47,11 @@ public String getExcludeSuffix(){
             return r.get("excludeSuffix").toString();
         } catch (FileNotFoundException e) {
             e.printStackTrace();
-            return "css|jpeg|gif|jpg|png|pdf|rar|zip|docx|doc|svg|jpeg|ico|woff|woff2|ttf|otf";
+            return Config.excludeSuffix;
         }
     }
 
-    public String getConfigPath(){
+    public static String getConfigPath(){
         try {
             InputStream inorder = new FileInputStream(SettingPath);
             Map<String,Object> r;
@@ -65,7 +68,7 @@ public void setExcludeSuffix(@NotNull String excludeSuffix){
         r.put("excludeSuffix", excludeSuffix);
         r.put("configPath", getConfigPath());
         try{
-            Writer ws = new OutputStreamWriter(new FileOutputStream(SettingPath),"UTF-8");
+            Writer ws = new OutputStreamWriter(new FileOutputStream(SettingPath), StandardCharsets.UTF_8);
             yaml.dump(r, ws);
         }catch (Exception ex){
             ex.printStackTrace();
@@ -77,7 +80,7 @@ public void setConfigPath(@NotNull String filePath){
         r.put("configPath", filePath);
         r.put("excludeSuffix", getExcludeSuffix());
         try{
-            Writer ws = new OutputStreamWriter(new FileOutputStream(SettingPath),"UTF-8");
+            Writer ws = new OutputStreamWriter(new FileOutputStream(SettingPath), StandardCharsets.UTF_8);
             yaml.dump(r, ws);
         }catch (Exception ex){
             ex.printStackTrace();

diff --git a/src/main/java/burp/yaml/LoadRule.java b/src/main/java/burp/yaml/LoadRule.java
@@ -7,6 +7,7 @@
 import org.yaml.snakeyaml.nodes.Tag;
 
 import java.io.*;
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Map;
@@ -17,16 +18,15 @@
 
 public class LoadRule {
     private static String filePath = "Config.yml";
-    public LoadRule(String configfile){
-        filePath = configfile;
+    public LoadRule(String configFile){
+        filePath = configFile;
         init();
     }
 
     // 初始化配置
     public void init(){
-        File settingyaml = new File(filePath);
-        if (!(settingyaml.exists() && settingyaml.isFile())){
-            Map<String,Object[][]> r = new HashMap<>();
+        File yamlFile = new File(filePath);
+        if (!(yamlFile.exists() && yamlFile.isFile())){
             Rule rule = new Rule();
             rule.setLoaded(true);
             rule.setName("Email");
@@ -50,10 +50,10 @@ public void init(){
             representer.addClassTag(Config.class, Tag.MAP);
 
             Yaml yaml = new Yaml(new Constructor(),representer,dop);
-            LoadConfigFile loadfile = new LoadConfigFile();
-            File f = new File(loadfile.getConfigPath());
+            new LoadConfigFile();
+            File f = new File(LoadConfigFile.getConfigPath());
             try{
-                Writer ws = new OutputStreamWriter(new FileOutputStream(f),"UTF-8");
+                Writer ws = new OutputStreamWriter(new FileOutputStream(f), StandardCharsets.UTF_8);
                 yaml.dump(config,ws);
             }catch (Exception ex){
                 ex.printStackTrace();
@@ -65,7 +65,7 @@ public static Map<String,Object[][]> getConfig(){
         InputStream inorder = null;
         {
             try {
-                inorder = new FileInputStream(new File(filePath));
+                inorder = new FileInputStream(filePath);
             } catch (FileNotFoundException e) {
                 e.printStackTrace();
             }