From d82eb3d5415426f2371de280c3471e08b1fb1e94 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Thu, 17 Dec 2020 16:42:13 +0800
Subject: [PATCH] dump level

---
 README.md             |  2 +-
 addon/dumper.go       | 34 +++++++++++++++++++++++++++++++---
 cmd/mitmproxy/main.go |  8 +++++---
 3 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 7225b92..89b2a05 100644
--- a/README.md
+++ b/README.md
@@ -14,6 +14,6 @@
 - [x] websocket
 - [x] 插件机制
 - [x] 命令行参数控制 dump 至文件
-- [ ] dump level
+- [x] dump level
 - [ ] support get method with body
 - [ ] http2
diff --git a/addon/dumper.go b/addon/dumper.go
index 06d9daa..636ae2f 100644
--- a/addon/dumper.go
+++ b/addon/dumper.go
@@ -7,21 +7,28 @@ import (
 	"net/http"
 	"os"
 	"strings"
+	"unicode"
 
 	"github.com/lqqyt2423/go-mitmproxy/flow"
 )
 
 type Dumper struct {
 	Base
-	Out io.Writer
+	level int // 0: header 1: header + body
+	Out   io.Writer
 }
 
-func NewDumperWithFile(file string) *Dumper {
+func NewDumper(file string, level int) *Dumper {
 	out, err := os.OpenFile(file, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
 	if err != nil {
 		panic(err)
 	}
-	return &Dumper{Out: out}
+
+	if level != 0 && level != 1 {
+		level = 0
+	}
+
+	return &Dumper{Out: out, level: level}
 }
 
 func (d *Dumper) Requestheaders(f *flow.Flow) {
@@ -48,6 +55,11 @@ func (d *Dumper) Requestheaders(f *flow.Flow) {
 		}
 		buf.WriteString("\r\n")
 
+		if d.level == 1 && f.Request.Body != nil && len(f.Request.Body) > 0 && CanPrint(f.Request.Body) {
+			buf.Write(f.Request.Body)
+			buf.WriteString("\r\n\r\n")
+		}
+
 		if f.Response != nil {
 			fmt.Fprintf(buf, "%v %v %v\r\n", f.Request.Proto, f.Response.StatusCode, http.StatusText(f.Response.StatusCode))
 			err = f.Response.Header.WriteSubset(buf, nil)
@@ -55,11 +67,27 @@ func (d *Dumper) Requestheaders(f *flow.Flow) {
 				log.Error(err)
 			}
 			buf.WriteString("\r\n")
+
+			if d.level == 1 && f.Response.Body != nil && len(f.Response.Body) > 0 && CanPrint(f.Response.Body) {
+				buf.Write(f.Response.Body)
+				buf.WriteString("\r\n\r\n")
+			}
 		}
 
+		buf.WriteString("\r\n\r\n")
+
 		_, err = d.Out.Write(buf.Bytes())
 		if err != nil {
 			log.Error(err)
 		}
 	}()
 }
+
+func CanPrint(content []byte) bool {
+	for _, c := range string(content) {
+		if !unicode.IsPrint(c) && !unicode.IsSpace(c) {
+			return false
+		}
+	}
+	return true
+}
diff --git a/cmd/mitmproxy/main.go b/cmd/mitmproxy/main.go
index da31d6d..3e611d1 100644
--- a/cmd/mitmproxy/main.go
+++ b/cmd/mitmproxy/main.go
@@ -10,8 +10,9 @@ import (
 )
 
 type Config struct {
-	addr string
-	dump string // dump filename
+	addr      string
+	dump      string // dump filename
+	dumpLevel int    // dump level
 }
 
 func loadConfig() *Config {
@@ -19,6 +20,7 @@ func loadConfig() *Config {
 
 	flag.StringVar(&config.addr, "addr", ":9080", "proxy listen addr")
 	flag.StringVar(&config.dump, "dump", "", "dump filename")
+	flag.IntVar(&config.dumpLevel, "dump_level", 0, "dump level: 0 - header, 1 - header + body")
 	flag.Parse()
 
 	return config
@@ -45,7 +47,7 @@ func main() {
 	}
 
 	if config.dump != "" {
-		dumper := addon.NewDumperWithFile(config.dump)
+		dumper := addon.NewDumper(config.dump, config.dumpLevel)
 		p.AddAddon(dumper)
 	}