From 16a0445afc30bcfe808f09ee5b8ea2946b08ca09 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Tue, 9 Feb 2021 11:51:17 +0800
Subject: [PATCH] addon/web.go -> addon/web/web.go

---
 addon/{ => web}/web.go | 26 ++++++++++++++------------
 cmd/mitmproxy/main.go  |  3 ++-
 2 files changed, 16 insertions(+), 13 deletions(-)
 rename addon/{ => web}/web.go (91%)

diff --git a/addon/web.go b/addon/web/web.go
similarity index 91%
rename from addon/web.go
rename to addon/web/web.go
index 9a558bc..93acfaa 100644
--- a/addon/web.go
+++ b/addon/web/web.go
@@ -1,4 +1,4 @@
-package addon
+package web
 
 import (
 	"encoding/json"
@@ -7,14 +7,17 @@ import (
 	"text/template"
 
 	"github.com/gorilla/websocket"
+	"github.com/lqqyt2423/go-mitmproxy/addon"
 	"github.com/lqqyt2423/go-mitmproxy/flow"
-	"github.com/sirupsen/logrus"
+	_log "github.com/sirupsen/logrus"
 )
 
+var log = _log.WithField("at", "web addon")
+
 func (web *WebAddon) echo(w http.ResponseWriter, r *http.Request) {
 	c, err := web.upgrader.Upgrade(w, r, nil)
 	if err != nil {
-		web.log.Print("upgrade:", err)
+		log.Print("upgrade:", err)
 		return
 	}
 
@@ -27,13 +30,13 @@ func (web *WebAddon) echo(w http.ResponseWriter, r *http.Request) {
 	for {
 		mt, message, err := c.ReadMessage()
 		if err != nil {
-			web.log.Println("read:", err)
+			log.Println("read:", err)
 			break
 		}
-		web.log.Printf("recv: %s", message)
+		log.Printf("recv: %s", message)
 		err = c.WriteMessage(mt, message)
 		if err != nil {
-			web.log.Println("write:", err)
+			log.Println("write:", err)
 			break
 		}
 	}
@@ -44,12 +47,11 @@ func home(w http.ResponseWriter, r *http.Request) {
 }
 
 type WebAddon struct {
-	Base
+	addon.Base
 	addr      string
 	upgrader  *websocket.Upgrader
 	serverMux *http.ServeMux
 	server    *http.Server
-	log       *logrus.Entry
 
 	conns   []*websocket.Conn
 	connsMu sync.RWMutex
@@ -77,13 +79,13 @@ func NewWebAddon() *WebAddon {
 	web.serverMux.HandleFunc("/", home)
 
 	web.server = &http.Server{Addr: web.addr, Handler: web.serverMux}
-	web.log = log.WithField("in", "WebAddon")
+	log = log.WithField("in", "WebAddon")
 	web.conns = make([]*websocket.Conn, 0)
 
 	go func() {
-		web.log.Infof("server start listen at %v\n", web.addr)
+		log.Infof("server start listen at %v\n", web.addr)
 		err := web.server.ListenAndServe()
-		web.log.Error(err)
+		log.Error(err)
 	}()
 
 	return web
@@ -125,7 +127,7 @@ func (web *WebAddon) sendFlow(on string, f *flow.Flow) {
 	msg := newMessage(on, f)
 	b, err := json.Marshal(msg)
 	if err != nil {
-		web.log.Error(err)
+		log.Error(err)
 		return
 	}
 	for _, c := range conns {
diff --git a/cmd/mitmproxy/main.go b/cmd/mitmproxy/main.go
index 28a9ad4..5af9482 100644
--- a/cmd/mitmproxy/main.go
+++ b/cmd/mitmproxy/main.go
@@ -5,6 +5,7 @@ import (
 	"os"
 
 	"github.com/lqqyt2423/go-mitmproxy/addon"
+	"github.com/lqqyt2423/go-mitmproxy/addon/web"
 	"github.com/lqqyt2423/go-mitmproxy/proxy"
 	log "github.com/sirupsen/logrus"
 )
@@ -51,7 +52,7 @@ func main() {
 		p.AddAddon(dumper)
 	}
 
-	p.AddAddon(addon.NewWebAddon())
+	p.AddAddon(web.NewWebAddon())
 
 	log.Fatal(p.Start())
 }