Highlight header or body

Author: cak

src/main/kotlin/ReflectHighlight.kt

@@ -4,15 +4,21 @@ import org.parosproxy.paros.network.HttpMessage
 import org.parosproxy.paros.view.View
 import org.zaproxy.zap.extension.search.SearchMatch
 
-class ReflectHighlight(private val reqRes: HttpMessage, private val tokens: List<String>) {
+class ReflectHighlight(
+    private val reqRes: HttpMessage,
+    private val headerTokens: List<String>,
+    private val bodyTokens: List<String>
+) {
 
-    private val searchMatch: SearchMatch? by lazy {
-        searchMatcher(tokens)
+    private val searchMatch: Pair<SearchMatch?, SearchMatch?> by lazy {
+        Pair(searchMatcher(headerTokens), searchMatcher(bodyTokens))
     }
 
     fun highlight() {
-        searchMatch?.let {
-            View.getSingleton().responsePanel.highlightBody(it)
+        if (searchMatch.second != null) {
+            View.getSingleton().responsePanel.highlightBody(searchMatch.second)
+        } else if (searchMatch.first != null) {
+            View.getSingleton().responsePanel.highlightHeader(searchMatch.first)
         }
     }