From 043f32ccb89fa512ea910a04eb5c12f7a63e488c Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Tue, 17 Jan 2023 16:15:47 +0800
Subject: [PATCH] web flow filter support scope

---
 web/client/src/lib/flow.ts | 107 +++++++++++++++++++++++++++++++++----
 1 file changed, 98 insertions(+), 9 deletions(-)

diff --git a/web/client/src/lib/flow.ts b/web/client/src/lib/flow.ts
index f6ddfc4..7e9b99c 100644
--- a/web/client/src/lib/flow.ts
+++ b/web/client/src/lib/flow.ts
@@ -262,13 +262,26 @@ export class Flow {
   }
 }
 
+const FLOW_FILTER_SCOPES = ['url', 'method', 'code', 'header', 'reqheader', 'resheader', 'body', 'reqbody', 'resbody', 'all'] as const
+type FlowFilterScope = typeof FLOW_FILTER_SCOPES[number]
+
 class FlowFilter {
   private keyword: string | RegExp | undefined
+  private scope: FlowFilterScope = 'url'
 
   constructor(text: string) {
     text = text.trim()
     if (!text) return
 
+    for (const scope of FLOW_FILTER_SCOPES) {
+      if (text.startsWith(`${scope}:`)) {
+        this.scope = scope
+        text = text.replace(`${scope}:`, '').trim()
+        break
+      }
+    }
+    if (!text) return
+
     // regexp
     if (text.startsWith('/') && (text.endsWith('/') || text.endsWith('/i'))) {
       let flags: string | undefined
@@ -286,10 +299,90 @@ class FlowFilter {
     }
   }
 
-  match(flow: Flow): boolean {
+  public match(flow: Flow): boolean {
+    switch (this.scope) {
+    case 'url':
+      return this.matchUrl(flow)
+    case 'method':
+      return this.matchMethod(flow)
+    case 'code':
+      return this.matchCode(flow)
+    case 'reqheader':
+      return this.matchReqHeader(flow)
+    case 'resheader':
+      return this.matchResHeader(flow)
+    case 'header':
+      return this.matchHeader(flow)
+    case 'reqbody':
+      return this.matchReqBody(flow)
+    case 'resbody':
+      return this.matchResBody(flow)
+    case 'body':
+      return this.matchBody(flow)
+    case 'all':
+      return this.matchAll(flow)
+    default:
+      throw new Error(`invalid scope ${this.scope}`)
+    }
+  }
+
+  private matchUrl(flow: Flow): boolean {
+    return this.matchKeyword(flow.request.url)
+  }
+
+  private matchMethod(flow: Flow): boolean {
+    return this.matchKeyword(flow.request.method) || this.matchKeyword(flow.request.method.toLowerCase())
+  }
+
+  private matchCode(flow: Flow): boolean {
+    if (!flow.response) return false
+    return this.matchKeyword(flow.response.statusCode.toString())
+  }
+
+  private _matchHeader(header: Header): boolean {
+    return Object.entries(header).some(([key, vals]) => {
+      return [key].concat(vals).some(text => this.matchKeyword(text))
+    })
+  }
+
+  private matchReqHeader(flow: Flow): boolean {
+    return this._matchHeader(flow.request.header)
+  }
+
+  private matchResHeader(flow: Flow): boolean {
+    if (!flow.response) return false
+    return this._matchHeader(flow.response.header)
+  }
+
+  private matchHeader(flow: Flow): boolean {
+    return this.matchReqHeader(flow) || this.matchResHeader(flow)
+  }
+
+  private matchReqBody(flow: Flow): boolean {
+    const body = flow.requestBody()
+    if (!body) return false
+    return this.matchKeyword(body)
+  }
+
+  private matchResBody(flow: Flow): boolean {
+    const body = flow.responseBody()
+    if (!body) return false
+    return this.matchKeyword(body)
+  }
+
+  private matchBody(flow: Flow): boolean {
+    return this.matchReqBody(flow) || this.matchResBody(flow)
+  }
+
+  private matchAll(flow: Flow): boolean {
+    return this.matchUrl(flow) || this.matchMethod(flow) || this.matchHeader(flow) || this.matchBody(flow)
+  }
+
+  private matchKeyword(text: string): boolean {
     if (!this.keyword) return true
-    if (this.keyword instanceof RegExp) return this.keyword.test(flow.request.url)
-    return flow.request.url.includes(this.keyword)
+    if (!text) return false
+    if (this.keyword instanceof RegExp) return this.keyword.test(text)
+    return text.includes(this.keyword)
   }
 }
 
@@ -330,10 +423,6 @@ export class FlowManager {
     return this._map.get(id)
   }
 
-  changeFilter(text: string) {
-    this.flowFilter = new FlowFilter(text)
-  }
-
   changeFilterLazy(text: string, callback: (err: any) => void) {
     if (this.filterTimer) {
       clearTimeout(this.filterTimer)
@@ -342,10 +431,10 @@ export class FlowManager {
 
     this.filterTimer = setTimeout(() => {
       try {
-        this.changeFilter(text)
+        this.flowFilter = new FlowFilter(text)
         callback(null)
       } catch (err) {
-        this.changeFilter('')
+        this.flowFilter = undefined
         callback(err)
       }
     }, 300) as any