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{} {