Version: 2.1.1 Update

src/main/java/burp/BurpExtender.java

@@ -6,10 +6,9 @@ import burp.ui.MainUI;
 import javax.swing.*;
 import java.awt.*;
 import java.nio.charset.StandardCharsets;
-import java.util.Arrays;
-import java.util.List;
 import java.io.PrintWriter;
-import java.util.Map;
+import java.util.ArrayList;
+import java.util.List;
 
 /*
  * @author EvilChen
@@ -20,11 +19,9 @@ public class BurpExtender implements IBurpExtender, IHttpListener, IMessageEdito
     private static PrintWriter stdout;
     private IBurpExtenderCallbacks callbacks;
     private static IExtensionHelpers helpers;
-    MatchHTTP mh = new MatchHTTP();
-    ExtractContent ec = new ExtractContent();
-    DoAction da = new DoAction();
     GetColorKey gck = new GetColorKey();
     UpgradeColor uc = new UpgradeColor();
+    ProcessMessage pm = new ProcessMessage();
 
     @Override
     public void registerExtenderCallbacks(final IBurpExtenderCallbacks callbacks)
@@ -32,7 +29,7 @@ public class BurpExtender implements IBurpExtender, IHttpListener, IMessageEdito
         this.callbacks = callbacks;
         BurpExtender.helpers = callbacks.getHelpers();
 
-        String version = "2.1";
+        String version = "2.1.1";
         callbacks.setExtensionName(String.format("HaE (%s) - Highlighter and Extractor", version));
         // 定义输出
         stdout = new PrintWriter(callbacks.getStdout(), true);
@@ -66,53 +63,29 @@ public class BurpExtender implements IBurpExtender, IHttpListener, IMessageEdito
     public void processHttpMessage(int toolFlag, boolean messageIsRequest, IHttpRequestResponse messageInfo) {
         // 判断是否是响应，且该代码作用域为：REPEATER、INTRUDER、PROXY（分别对应toolFlag 64、32、4）
         if (toolFlag == 64 || toolFlag == 32 || toolFlag == 4) {
-            Map<String, Map<String, Object>> obj;
-            // 流量清洗
-            String urlString = helpers.analyzeRequest(messageInfo.getHttpService(), messageInfo.getRequest()).getUrl().toString();
-            urlString = urlString.indexOf("?") > 0 ? urlString.substring(0, urlString.indexOf("?")) : urlString;
-
-            // 正则判断
-            if (mh.matchSuffix(urlString)) {
-                return;
-            }
-
+            byte[] content;
             if (messageIsRequest) {
-                byte[] byteRequest = messageInfo.getRequest();
-                // 获取报文头
-                List<String> requestTmpHeaders = helpers.analyzeRequest(messageInfo.getHttpService(), byteRequest).getHeaders();
-                String requestHeaders = String.join("\n", requestTmpHeaders);
-
-                // 获取报文主体
-                int requestBodyOffset = helpers.analyzeRequest(messageInfo.getHttpService(), byteRequest).getBodyOffset();
-                byte[] requestBody = Arrays.copyOfRange(byteRequest, requestBodyOffset, byteRequest.length);
-
-                obj = ec.matchRegex(byteRequest, requestHeaders, requestBody, "request");
+                content = messageInfo.getRequest();
             } else {
-                byte[] byteResponse = messageInfo.getResponse();
-
-                // 获取报文头
-                List<String> responseTmpHeaders = helpers.analyzeRequest(messageInfo.getHttpService(), byteResponse).getHeaders();
-                String responseHeaders = String.join("\n", responseTmpHeaders);
-
-                // 获取报文主体
-                int responseBodyOffset = helpers.analyzeResponse(byteResponse).getBodyOffset();
-                byte[] responseBody = Arrays.copyOfRange(byteResponse, responseBodyOffset, byteResponse.length);
-
-                obj = ec.matchRegex(byteResponse, responseHeaders, responseBody, "response");
+                content = messageInfo.getResponse();
             }
-
-            List<List<String>> resultList = da.highlightAndComment(obj);
-            List<String> colorList = resultList.get(0);
-            stdout.println(colorList);
-            List<String> commentList = resultList.get(1);
-            if (colorList.size() != 0) {
-                String color = uc.getEndColor(gck.getColorKeys(colorList));
-                messageInfo.setHighlight(color);
-            }
-
-            if (commentList.size() != 0) {
+            String c = new String(content, StandardCharsets.UTF_8).intern();
+            List<String> result = pm.processMessageByContent(helpers, messageInfo.getHttpService(), content, messageIsRequest, true);
+            if (result != null && !result.isEmpty() && result.size() > 0) {
+                String originalColor = messageInfo.getHighlight();
                 String originalComment = messageInfo.getComment();
-                messageInfo.setComment(String.join(", ", commentList));
+                List<String> colorList = new ArrayList<>();
+                if (originalColor != null) {
+                    colorList.add(originalColor);
+                }
+                colorList.add(result.get(0));
+                String color = uc.getEndColor(gck.getColorKeys(colorList));
+
+                messageInfo.setHighlight(color);
+                String addComment = String.join(", ", result.get(1));
+                String resComment = originalComment != null ? String.format("%s, %s", originalComment, addComment) : addComment;
+
+                messageInfo.setComment(resComment);
             }
         }
 
@@ -143,52 +116,20 @@ public class BurpExtender implements IBurpExtender, IHttpListener, IMessageEdito
 
         @Override
         public boolean isEnabled(byte[] content, boolean isRequest) {
-            Map<String, Map<String, Object>> obj;
+            String c = new String(content, StandardCharsets.UTF_8).intern();
+            List<String> result = pm.processMessageByContent(helpers, controller.getHttpService(), content, isRequest, false);
 
-            if (isRequest) {
-                try {
-                    // 流量清洗
-                    String urlString = helpers.analyzeRequest(controller.getHttpService(), controller.getRequest()).getUrl().toString();
-                    urlString = urlString.indexOf("?") > 0 ? urlString.substring(0, urlString.indexOf("?")) : urlString;
-                    // 正则判断
-                    if (mh.matchSuffix(urlString)) {
-                        return false;
-                    }
-                } catch (Exception e) {
-                    return false;
-                }
-                IRequestInfo iRequestInfo = helpers.analyzeRequest(controller.getHttpService(), content);
-
-                // 获取报文头
-                List<String> requestTmpHeaders = iRequestInfo.getHeaders();
-                String requestHeaders = String.join("\n", requestTmpHeaders);
-                // 获取报文主体
-                int requestBodyOffset = iRequestInfo.getBodyOffset();
-                byte[] requestBody = Arrays.copyOfRange(content, requestBodyOffset, content.length);
-
-                obj = ec.matchRegex(content, requestHeaders, requestBody, "request");
-                if (obj.size() > 0) {
-                    String result = da.extractString(obj);
-                    extractRequestContent = result.getBytes();
-                    return true;
-                }
-            } else {
-                IResponseInfo iResponseInfo = helpers.analyzeResponse(content);
-                // 获取报文头
-                List<String> responseTmpHeaders = iResponseInfo.getHeaders();
-                String responseHeaders = String.join("\n", responseTmpHeaders);
-                // 获取报文主体
-                int responseBodyOffset = iResponseInfo.getBodyOffset();
-                byte[] responseBody = Arrays.copyOfRange(content, responseBodyOffset, content.length);
-
-                obj = ec.matchRegex(content, responseHeaders, responseBody, "response");
-                if (obj.size() > 0) {
-                    String result = da.extractString(obj);
-                    extractResponseContent = result.getBytes();
-                    return true;
+            if (result != null && !result.isEmpty()) {
+                if (isRequest) {
+                    extractRequestContent = result.get(0).getBytes();
+                } else {
+                    extractResponseContent = result.get(0).getBytes();
                 }
+                return true;
             }
+
             return false;
+
         }
 
         @Override

src/main/java/burp/action/ProcessMessage.java

@@ -0,0 +1,85 @@
+package burp.action;
+
+import burp.IExtensionHelpers;
+import burp.IHttpService;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+
+public class ProcessMessage {
+    MatchHTTP mh = new MatchHTTP();
+    ExtractContent ec = new ExtractContent();
+    DoAction da = new DoAction();
+    GetColorKey gck = new GetColorKey();
+    UpgradeColor uc = new UpgradeColor();
+
+    public List<String> processMessageByContent(IExtensionHelpers helpers, IHttpService httpService, byte[] content, boolean isRequest, boolean messageInfo) {
+        List<String> result = new ArrayList<>();;
+        Map<String, Map<String, Object>> obj;
+
+        if (isRequest) {
+            try {
+                // 流量清洗
+                String urlString = helpers.analyzeRequest(httpService, content).getUrl().toString();
+                urlString = urlString.indexOf("?") > 0 ? urlString.substring(0, urlString.indexOf("?")) : urlString;
+
+                // 正则判断
+                if (mh.matchSuffix(urlString)) {
+                    return result;
+                }
+            } catch (Exception e) {
+                return result;
+            }
+            // 获取报文头
+            List<String> requestTmpHeaders = helpers.analyzeRequest(httpService, content).getHeaders();
+            String requestHeaders = String.join("\n", requestTmpHeaders);
+
+            // 获取报文主体
+            int requestBodyOffset = helpers.analyzeRequest(httpService, content).getBodyOffset();
+            byte[] requestBody = Arrays.copyOfRange(content, requestBodyOffset, content.length);
+
+            obj = ec.matchRegex(content, requestHeaders, requestBody, "request");
+        } else {
+            try {
+                // 流量清洗
+                String inferredMimeType = String.format("hae.%s", helpers.analyzeResponse(content).getInferredMimeType().toLowerCase());
+                String statedMimeType = String.format("hae.%s", helpers.analyzeResponse(content).getStatedMimeType().toLowerCase());
+                // 正则判断
+                if (mh.matchSuffix(statedMimeType) || mh.matchSuffix(inferredMimeType)) {
+                    return result;
+                }
+            } catch (Exception e) {
+                return result;
+            }
+            // 获取报文头
+            List<String> responseTmpHeaders = helpers.analyzeResponse(content).getHeaders();
+            String responseHeaders = String.join("\n", responseTmpHeaders);
+
+            // 获取报文主体
+            int responseBodyOffset = helpers.analyzeResponse(content).getBodyOffset();
+            byte[] responseBody = Arrays.copyOfRange(content, responseBodyOffset, content.length);
+
+            obj = ec.matchRegex(content, responseHeaders, responseBody, "response");
+        }
+
+        if (messageInfo) {
+            List<List<String>> resultList = da.highlightAndComment(obj);
+            List<String> colorList = resultList.get(0);
+            List<String> commentList = resultList.get(1);
+            if (colorList.size() != 0 && commentList.size() != 0) {
+                String color = uc.getEndColor(gck.getColorKeys(colorList));
+                result.add(color);
+                result.add(String.join(", ", commentList));
+            }
+        } else {
+            if (obj.size() > 0) {
+                result.add(da.extractString(obj));
+            }
+        }
+
+        return result;
+
+    }
+}