增加改变请求头、请求cookie、响应头、响应cookie的能力

Leviathangk · May 11, 2023 · 8475706 · 8475706
1 parent d6a3f3c
commit 8475706
Show file tree

Hide file tree

Showing 5 changed files with 211 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -11,13 +11,15 @@ go get github.com/Leviathangk/go-mitmtools@latest
 
 # Handler
 
-主要有以下 Handler
+主要分为请求处理、响应处理
 
 ## req
 
 请求 handler
 
 - ShowReq：打印请求
+- ChangeHeader：修改请求头（注意大小写）
+- ChangeCookie：修改请求 Cookie（注意大小写）
 
 ## resp
 
@@ -34,6 +36,8 @@ go get github.com/Leviathangk/go-mitmtools@latest
 - FindContent：输出含有指定字符的 url
 - FindCookie：输出含有指定 响应 Cookie 的 url：匹配的就是 document.cookie 后的那部分
 - FindHeader：输出含有指定 响应头 的 url：匹配的是响应头的 key
+- ChangeHeader：修改响应头
+- ChangeCookie：修改响应 Cookie
 
 # 案例
 
@@ -133,6 +137,30 @@ func main() {
 		Pattern:    "^https://www.baidu.com/$",
 		KeyPattern: []string{"Bdqid", "Set-Cookie"},
 	})
+	
+	// 修改响应头
+	opts.AddHandler(&resp.ChangeHeader{
+		Pattern: "^https://www.baidu.com/$",
+		Header:  map[string][]string{"Bdqid": {"baidu"}},
+	})
+	
+	// 修改响应 cookie
+	opts.AddHandler(&resp.ChangeCookie{
+		Pattern: "^https://www.baidu.com/$",
+		Cookie:  map[string]string{"H_PS_PSSID": "baidu"},
+	})
+	
+	// 修改请求头
+	opts.AddHandler(&req.ChangeHeader{
+		Pattern: "^http://127.0.0.1:8877/headerTest$",
+		Header:  map[string][]string{"X": {"qiandu"}},
+	})
+	
+	// 修改请求 cookie
+	opts.AddHandler(&req.ChangeCookie{
+		Pattern: "^http://127.0.0.1:8877/cookieTest$",
+		Cookie:  map[string]string{"x": "qiandu"},
+	})
 
 	glog.Fatal(mitmtools.Start(opts))
 }

diff --git a/demo.go b/demo.go
@@ -94,5 +94,29 @@ func main() {
 		KeyPattern: []string{"Bdqid", "Set-Cookie"},
 	})
 
+	// 修改响应头
+	opts.AddHandler(&resp.ChangeHeader{
+		Pattern: "^https://www.baidu.com/$",
+		Header:  map[string][]string{"Bdqid": {"baidu"}},
+	})
+
+	// 修改响应 cookie
+	opts.AddHandler(&resp.ChangeCookie{
+		Pattern: "^https://www.baidu.com/$",
+		Cookie:  map[string]string{"H_PS_PSSID": "baidu"},
+	})
+
+	// 修改请求头
+	opts.AddHandler(&req.ChangeHeader{
+		Pattern: "^http://127.0.0.1:8877/headerTest$",
+		Header:  map[string][]string{"X": {"qiandu"}},
+	})
+
+	// 修改请求 cookie
+	opts.AddHandler(&req.ChangeCookie{
+		Pattern: "^http://127.0.0.1:8877/cookieTest$",
+		Cookie:  map[string]string{"x": "qiandu"},
+	})
+
 	glog.Fatal(mitmtools.Start(opts))
 }
diff --git a/mitmtools/handler/req/change.go b/mitmtools/handler/req/change.go
@@ -0,0 +1,77 @@
+package req
+
+import (
+	"fmt"
+	"github.com/Leviathangk/go-glog/glog"
+	"github.com/Leviathangk/go-mitmtools/mitmtools/handler"
+	"github.com/lqqyt2423/go-mitmproxy/proxy"
+	"strings"
+)
+
+type ChangeHeader struct {
+	handler.BaseHandler
+	Pattern string              // url 匹配规则
+	Header  map[string][]string // 需要替换的头
+}
+
+func (r *ChangeHeader) Requestheaders(f *proxy.Flow) {
+	if handler.IsMatch(r.Pattern, f.Request.URL.String()) {
+		for key, value := range r.Header {
+			if _, ok := f.Request.Header[key]; ok {
+				f.Request.Header[key] = value
+
+				if handler.ShowLog {
+					glog.Debugf("ChangeHeader 正在替换指定请求头：%s -> %v\n", key, value)
+				}
+			}
+		}
+	}
+}
+
+// Check 检查是否符合启动要求
+func (r *ChangeHeader) Check() error {
+
+	if r.Header == nil {
+		return fmt.Errorf("ChangeHeader 未接收到需要替换的请求头！")
+	}
+
+	return nil
+}
+
+type ChangeCookie struct {
+	handler.BaseHandler
+	Pattern string            // url 匹配规则
+	Cookie  map[string]string // 需要查找的内容
+}
+
+func (fin *ChangeCookie) Requestheaders(f *proxy.Flow) {
+	if handler.IsMatch(fin.Pattern, f.Request.URL.String()) {
+		for key, value := range f.Request.Header {
+			if strings.ToLower(key) == "cookie" {
+				for keyIndex, cookie := range value {
+					for newK, newV := range fin.Cookie {
+						if handler.IsMatch(newK, cookie) {
+							newCookie := newK + "=" + newV
+							f.Request.Header[key][keyIndex] = newCookie
+
+							if handler.ShowLog {
+								glog.Debugf("ChangeCookie 已查找到：%s -> %s -> %s\n", newK, newCookie, f.Request.URL)
+							}
+						}
+					}
+				}
+				break
+			}
+		}
+	}
+}
+
+// Check 检查是否符合启动要求
+func (fin *ChangeCookie) Check() error {
+
+	if len(fin.Cookie) == 0 {
+		return fmt.Errorf("ChangeCookie 未接收到需要查找的内容！")
+	}
+
+	return nil
+}
diff --git a/mitmtools/handler/resp/change.go b/mitmtools/handler/resp/change.go
@@ -0,0 +1,80 @@
+package resp
+
+import (
+	"fmt"
+	"github.com/Leviathangk/go-glog/glog"
+	"github.com/Leviathangk/go-mitmtools/mitmtools/handler"
+	"github.com/lqqyt2423/go-mitmproxy/proxy"
+	"strings"
+)
+
+type ChangeHeader struct {
+	handler.BaseHandler
+	Pattern string              // url 匹配规则
+	Header  map[string][]string // 需要替换的头
+}
+
+func (r *ChangeHeader) Responseheaders(f *proxy.Flow) {
+	if handler.IsMatch(r.Pattern, f.Request.URL.String()) {
+		for key, value := range r.Header {
+			if _, ok := f.Response.Header[key]; ok {
+				f.Response.Header[key] = value
+
+				if handler.ShowLog {
+					glog.Debugf("ChangeHeader 正在替换指定请求头：%s -> %v\n", key, value)
+				}
+			}
+		}
+	}
+}
+
+// Check 检查是否符合启动要求
+func (r *ChangeHeader) Check() error {
+
+	if r.Header == nil {
+		return fmt.Errorf("ChangeHeader 未接收到需要替换的请求头！")
+	}
+
+	return nil
+}
+
+type ChangeCookie struct {
+	handler.BaseHandler
+	Pattern string            // url 匹配规则
+	Cookie  map[string]string // 需要查找的内容
+}
+
+func (fin *ChangeCookie) Response(f *proxy.Flow) {
+	if handler.IsMatch(fin.Pattern, f.Request.URL.String()) {
+		for key, value := range f.Response.Header {
+			if strings.ToLower(key) == "set-cookie" {
+				for keyIndex, cookie := range value {
+					for newK, newV := range fin.Cookie {
+						if handler.IsMatch(newK, cookie) {
+							cookieSlice := strings.Split(cookie, ";")
+
+							cookie = strings.ReplaceAll(cookie, cookieSlice[0], strings.Split(cookieSlice[0], "=")[0]+"="+newV)
+
+							f.Response.Header[key][keyIndex] = cookie
+
+							if handler.ShowLog {
+								glog.Debugf("ChangeCookie 已查找到：%s -> %s -> %s\n", newK, cookie, f.Request.URL)
+							}
+						}
+					}
+				}
+				break
+			}
+		}
+	}
+}
+
+// Check 检查是否符合启动要求
+func (fin *ChangeCookie) Check() error {
+
+	if len(fin.Cookie) == 0 {
+		return fmt.Errorf("ChangeCookie 未接收到需要查找的内容！")
+	}
+
+	return nil
+}
diff --git a/mitmtools/handler/resp/header.go b/mitmtools/handler/resp/header.go
@@ -17,7 +17,7 @@ type AddHeader struct {
 func (a *AddHeader) Responseheaders(f *proxy.Flow) {
 	if handler.IsMatch(a.Pattern, f.Request.URL.String()) {
 		for key, value := range a.Header {
-			f.Response.Header.Add(key, value)
+			f.Response.Header.Set(key, value)
 
 			if handler.ShowLog {
 				glog.Debugf("AddHeader 正在添加指定请求头：%s -> %s\n", key, value)