web addon todo

4 years ago · d352d25316
parent d8736ee98a
commit d352d25316
5 changed files with 99 additions and 6 deletions
--- 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(`
 <!DOCTYPE html>
 <html>
--- 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())
 }
--- 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{} {
--- 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
 )
--- 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=