package web
import (
"encoding/json"
"net/http"
"sync"
"text/template"
"github.com/gorilla/websocket"
"github.com/lqqyt2423/go-mitmproxy/addon"
"github.com/lqqyt2423/go-mitmproxy/flow"
_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 {
log.Print("upgrade:", err)
return
}
web.addConn(c)
defer func() {
web.removeConn(c)
c.Close()
}()
for {
mt, message, err := c.ReadMessage()
if err != nil {
log.Println("read:", err)
break
}
log.Printf("recv: %s", message)
err = c.WriteMessage(mt, message)
if err != nil {
log.Println("write:", err)
break
}
}
}
func home(w http.ResponseWriter, r *http.Request) {
homeTemplate.Execute(w, "ws://"+r.Host+"/echo")
}
type WebAddon struct {
addon.Base
addr string
upgrader *websocket.Upgrader
serverMux *http.ServeMux
server *http.Server
conns []*websocket.Conn
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"
web.upgrader = &websocket.Upgrader{}
web.serverMux = new(http.ServeMux)
web.serverMux.HandleFunc("/echo", web.echo)
web.serverMux.HandleFunc("/", home)
web.server = &http.Server{Addr: web.addr, Handler: web.serverMux}
log = log.WithField("in", "WebAddon")
web.conns = make([]*websocket.Conn, 0)
go func() {
log.Infof("server start listen at %v\n", web.addr)
err := web.server.ListenAndServe()
log.Error(err)
}()
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) sendFlow(on string, f *flow.Flow) {
web.connsMu.RLock()
conns := web.conns
web.connsMu.RUnlock()
if len(conns) == 0 {
return
}
msg := newMessage(on, f)
b, err := json.Marshal(msg)
if err != nil {
log.Error(err)
return
}
for _, c := range conns {
c.WriteMessage(websocket.TextMessage, 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(`
Click "Open" to create a connection to the server,
"Send" to send a message to the server and "Close" to close the connection.
You can change the message and send multiple times.