diff --git a/addon/web.go b/addon/web.go index ddf50ff..a9a8bda 100644 --- a/addon/web.go +++ b/addon/web.go @@ -1,10 +1,14 @@ package addon import ( + "encoding/json" "net/http" + "sync" "text/template" "github.com/gorilla/websocket" + "github.com/lqqyt2423/go-mitmproxy/flow" + uuid "github.com/satori/go.uuid" "github.com/sirupsen/logrus" ) @@ -14,7 +18,13 @@ func (web *WebAddon) echo(w http.ResponseWriter, r *http.Request) { web.log.Print("upgrade:", err) return } - defer c.Close() + + web.addConn(c) + defer func() { + web.removeConn(c) + c.Close() + }() + for { mt, message, err := c.ReadMessage() if err != nil { @@ -41,6 +51,9 @@ type WebAddon struct { serverMux *http.ServeMux server *http.Server log *logrus.Entry + + conns []*websocket.Conn + connsMu sync.RWMutex } func NewWebAddon() *WebAddon { @@ -53,8 +66,8 @@ func NewWebAddon() *WebAddon { web.serverMux.HandleFunc("/", home) web.server = &http.Server{Addr: web.addr, Handler: web.serverMux} - web.log = log.WithField("in", "WebAddon") + web.conns = make([]*websocket.Conn, 0) go func() { web.log.Infof("server start listen at %v\n", web.addr) @@ -65,6 +78,53 @@ func NewWebAddon() *WebAddon { return web } +func (web *WebAddon) addConn(c *websocket.Conn) { + web.connsMu.Lock() + web.conns = append(web.conns, c) + web.connsMu.Unlock() +} + +func (web *WebAddon) removeConn(conn *websocket.Conn) { + web.connsMu.Lock() + defer web.connsMu.Unlock() + + index := -1 + for i, c := range web.conns { + if conn == c { + index = i + break + } + } + + if index == -1 { + return + } + web.conns = append(web.conns[:index], web.conns[index+1:]...) +} + +func (web *WebAddon) Request(f *flow.Flow) { + b, err := json.Marshal(f) + if err != nil { + web.log.Error(err) + return + } + + id := uuid.NewV4() + f.State["id"] = id + + web.log.Infof("id: %s, request: %s\n", id, b) +} + +func (web *WebAddon) Response(f *flow.Flow) { + b, err := json.Marshal(f) + if err != nil { + web.log.Error(err) + return + } + + web.log.Infof("id: %s, response: %s\n", f.State["id"], b) +} + var homeTemplate = template.Must(template.New("").Parse(` diff --git a/cmd/mitmproxy/main.go b/cmd/mitmproxy/main.go index 3e611d1..28a9ad4 100644 --- a/cmd/mitmproxy/main.go +++ b/cmd/mitmproxy/main.go @@ -51,5 +51,7 @@ func main() { p.AddAddon(dumper) } + p.AddAddon(addon.NewWebAddon()) + log.Fatal(p.Start()) } diff --git a/flow/flow.go b/flow/flow.go index 2c4cb7d..7576cd6 100644 --- a/flow/flow.go +++ b/flow/flow.go @@ -1,6 +1,7 @@ package flow import ( + "encoding/json" "net/http" "net/url" @@ -19,6 +20,16 @@ type Request struct { raw *http.Request } +func (req *Request) MarshalJSON() ([]byte, error) { + r := make(map[string]interface{}) + r["method"] = req.Method + r["url"] = req.URL.String() + r["proto"] = req.Proto + r["header"] = req.Header + r["body"] = req.Body + return json.Marshal(r) +} + func NewRequest(req *http.Request) *Request { return &Request{ Method: req.Method, @@ -34,9 +45,9 @@ func (r *Request) Raw() *http.Request { } type Response struct { - StatusCode int - Header http.Header - Body []byte + StatusCode int `json:"statusCode"` + Header http.Header `json:"header"` + Body []byte `json:"body"` decodedBody []byte decoded bool // decoded reports whether the response was sent compressed but was decoded to decodedBody. @@ -51,10 +62,19 @@ type Flow struct { // 如果为 true，则不缓冲 Request.Body 和 Response.Body，且不进入之后的 Addon.Request 和 Addon.Response Stream bool done chan struct{} + + State map[string]interface{} // Can add value by addon +} + +func (f *Flow) MarshalJSON() ([]byte, error) { + j := make(map[string]interface{}) + j["request"] = f.Request + j["response"] = f.Response + return json.Marshal(j) } func NewFlow() *Flow { - return &Flow{done: make(chan struct{})} + return &Flow{done: make(chan struct{}), State: make(map[string]interface{})} } func (f *Flow) Done() <-chan struct{} { diff --git a/go.mod b/go.mod index 12a0f0c..c4bc7eb 100644 --- a/go.mod +++ b/go.mod @@ -7,5 +7,7 @@ require ( github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e github.com/gorilla/websocket v1.4.2 github.com/jordwest/mock-conn v0.0.0-20180617021051-4896c6bd1641 + github.com/satori/go.uuid v1.2.0 github.com/sirupsen/logrus v1.7.0 + gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect ) diff --git a/go.sum b/go.sum index 339e81f..22d4e6a 100644 --- a/go.sum +++ b/go.sum @@ -8,11 +8,20 @@ github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0U github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/jordwest/mock-conn v0.0.0-20180617021051-4896c6bd1641 h1:ChkB2s4mFDekyUUmbNE7qNhennP0rfqF2YZUOGxbhFk= github.com/jordwest/mock-conn v0.0.0-20180617021051-4896c6bd1641/go.mod h1:AJFEOPtj5Z5z3MAy+0uvjQAH02iRnQr6fnvuHYp/Jek= +github.com/kr/pretty v0.2.1 h1:Fmg33tUaq4/8ym9TJN1x7sLJnHVwhP33CNkpYV/7rwI= +github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/satori/go.uuid v1.2.0 h1:0uYX9dsZ2yD7q2RtLRtPSdGDWzjeM3TbMJP9utgA0ww= +github.com/satori/go.uuid v1.2.0/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0= github.com/sirupsen/logrus v1.7.0 h1:ShrD1U9pZB12TX0cVy0DtePoCH97K8EtX+mg7ZARUtM= github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0= github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w= github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= golang.org/x/sys v0.0.0-20191026070338-33540a1f6037 h1:YyJpGZS1sBuBCzLAR1VEpK193GlqGZbnPFnPV/5Rsb4= golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=