From 68a5284163b60af75ae57b5c7f54b1a196097151 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Mon, 8 Feb 2021 19:10:40 +0800
Subject: [PATCH] web addon

---
 addon/web.go | 44 ++++++++++++++++++++++++++++++--------------
 flow/flow.go |  9 ++++++++-
 2 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/addon/web.go b/addon/web.go
index a9a8bda..9a558bc 100644
--- a/addon/web.go
+++ b/addon/web.go
@@ -8,7 +8,6 @@ import (
 
 	"github.com/gorilla/websocket"
 	"github.com/lqqyt2423/go-mitmproxy/flow"
-	uuid "github.com/satori/go.uuid"
 	"github.com/sirupsen/logrus"
 )
 
@@ -56,6 +55,18 @@ type WebAddon struct {
 	connsMu sync.RWMutex
 }
 
+type message struct {
+	On   string     `json:"on"`
+	Flow *flow.Flow `json:"flow"`
+}
+
+func newMessage(on string, f *flow.Flow) *message {
+	return &message{
+		On:   on,
+		Flow: f,
+	}
+}
+
 func NewWebAddon() *WebAddon {
 	web := new(WebAddon)
 	web.addr = ":9081"
@@ -102,27 +113,32 @@ func (web *WebAddon) removeConn(conn *websocket.Conn) {
 	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)
+func (web *WebAddon) sendFlow(on string, f *flow.Flow) {
+	web.connsMu.RLock()
+	conns := web.conns
+	web.connsMu.RUnlock()
+
+	if len(conns) == 0 {
 		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)
+	msg := newMessage(on, f)
+	b, err := json.Marshal(msg)
 	if err != nil {
 		web.log.Error(err)
 		return
 	}
+	for _, c := range conns {
+		c.WriteMessage(websocket.TextMessage, b)
+	}
+}
 
-	web.log.Infof("id: %s, response: %s\n", f.State["id"], b)
+func (web *WebAddon) Request(f *flow.Flow) {
+	web.sendFlow("request", f)
+}
+
+func (web *WebAddon) Response(f *flow.Flow) {
+	web.sendFlow("response", f)
 }
 
 var homeTemplate = template.Must(template.New("").Parse(`
diff --git a/flow/flow.go b/flow/flow.go
index 7576cd6..db249ac 100644
--- a/flow/flow.go
+++ b/flow/flow.go
@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"net/url"
 
+	uuid "github.com/satori/go.uuid"
 	_log "github.com/sirupsen/logrus"
 )
 
@@ -63,18 +64,24 @@ type Flow struct {
 	Stream bool
 	done   chan struct{}
 
+	Id    uuid.UUID
 	State map[string]interface{} // Can add value by addon
 }
 
 func (f *Flow) MarshalJSON() ([]byte, error) {
 	j := make(map[string]interface{})
+	j["id"] = f.Id
 	j["request"] = f.Request
 	j["response"] = f.Response
 	return json.Marshal(j)
 }
 
 func NewFlow() *Flow {
-	return &Flow{done: make(chan struct{}), State: make(map[string]interface{})}
+	return &Flow{
+		done:  make(chan struct{}),
+		Id:    uuid.NewV4(),
+		State: make(map[string]interface{}),
+	}
 }
 
 func (f *Flow) Done() <-chan struct{} {