updated ipbl websocket feat

cmd/ipbl/ipbl.go

@@ -36,8 +36,8 @@ func main() {
 	if !cfg.Switchs.NoScanIP {
 		go models.ScanIP(&cfg)
 	}
-	//go zmqrouter.Init(&cfg)
+
-	go ws.Init(&cfg)
+	ws.Init(&cfg)
 	go func() { err = routers.RunServer(&ctx, &cfg) }()
 	if err != nil {
 		log.Fatalln(err)

src/routers/funcs.go

@@ -155,8 +155,12 @@ func RegisterRoutes(e *echo.Echo, ctx *context.Context, cfg *config.Config) {
 		return Result(c, err, disc)
 	})
 	e.File("/test.html", "/home/paul/test.html")
-	e.GET("/ws", func(c echo.Context) (err error) {
+	e.GET("/wsps", func(c echo.Context) (err error) {
-		ws.HandleWS(c, cfg)
+		ws.HandleWSPS(c, cfg)
+		return
+	})
+	e.GET("/wsrr", func(c echo.Context) (err error) {
+		ws.HandleWSRR(c, cfg)
 		return
 	})
 }

src/ws/init.go

@@ -0,0 +1,63 @@
+package ws
+
+import (
+	"time"
+
+	"git.paulbsd.com/paulbsd/ipbl/src/config"
+	"git.paulbsd.com/paulbsd/ipbl/src/models"
+	"golang.org/x/net/websocket"
+)
+
+var WebSocketChannel = make(chan models.APIEvent, 100)
+var WebSocketChannelDone = make(chan bool)
+
+var LISTENERS map[string]ConnectionInfo
+
+func Init(cfg *config.Config) {
+	LISTENERS = make(map[string]ConnectionInfo)
+}
+
+func WelcomeAgents(ws *websocket.Conn, welcome WSWelcome, t string) {
+	connectinfo, ok := LISTENERS[welcome.Hostname]
+
+	if !ok {
+		if t == "ps" {
+			connectinfo := ConnectionInfo{
+				ConnectionPS: ws,
+				InitDate:     time.Now(),
+			}
+			LISTENERS[welcome.Hostname] = connectinfo
+		} else if t == "rr" {
+			connectinfo := ConnectionInfo{
+				ConnectionRR: ws,
+				InitDate:     time.Now(),
+			}
+			LISTENERS[welcome.Hostname] = connectinfo
+		}
+	} else {
+		if t == "ps" {
+			connectinfo.ConnectionPS = ws
+		} else if t == "rr" {
+			connectinfo.ConnectionRR = ws
+		}
+	}
+}
+
+type ConnectionInfo struct {
+	ConnectionPS *websocket.Conn
+	ConnectionRR *websocket.Conn
+	InitDate     time.Time
+}
+
+type WSWelcome struct {
+	Hostname string
+}
+
+func gcConn(conn *websocket.Conn, t string) (err error) {
+	for index, value := range LISTENERS {
+		if value.ConnectionPS == conn {
+			delete(LISTENERS, index)
+		}
+	}
+	return err
+}

src/ws/pubsub.go

@@ -0,0 +1,51 @@
+package ws
+
+import (
+	"encoding/json"
+	"log"
+
+	"git.paulbsd.com/paulbsd/ipbl/src/config"
+	"github.com/labstack/echo/v4"
+	"golang.org/x/net/websocket"
+)
+
+func HandleWSPS(c echo.Context, cfg *config.Config) (err error) {
+	websocket.Handler(func(ws *websocket.Conn) {
+		defer ws.Close()
+
+		var welcome = WSWelcome{}
+
+		var msg []byte
+		err := websocket.Message.Receive(ws, &msg)
+		if err != nil {
+			log.Println(err)
+			return
+		}
+
+		err = json.Unmarshal(msg, &welcome)
+		if err != nil {
+			log.Println(err)
+		} else {
+			WelcomeAgents(ws, welcome, "ps")
+		}
+
+		/*for {
+			var apievent, ok = <-WebSocketChannel
+			if ok {
+				log.Println(apievent, ok)
+				data, err := json.Marshal(apievent)
+				if err != nil {
+					log.Println(err)
+				}
+				err = websocket.Message.Send(ws, data)
+				if err != nil {
+					gcConn(ws, "ps")
+					log.Println(LISTENERS)
+					return
+				}
+			}
+		}*/
+
+	}).ServeHTTP(c.Response(), c.Request())
+	return nil
+}

src/ws/reqrep.go

@@ -2,7 +2,6 @@ package ws
 
 import (
 	"encoding/json"
-	"fmt"
 	"log"
 
 	"git.paulbsd.com/paulbsd/ipbl/src/config"
@@ -11,29 +10,38 @@ import (
 	"golang.org/x/net/websocket"
 )
 
-var LISTENERS map[string]*websocket.Conn
+func HandleWSRR(c echo.Context, cfg *config.Config) error {
-var CHAN int
-
-func Init(cfg *config.Config) {
-	LISTENERS = make(map[string]*websocket.Conn)
-}
-
-func HandleWS(c echo.Context, cfg *config.Config) error {
 	websocket.Handler(func(ws *websocket.Conn) {
 		defer ws.Close()
 
+		var welcome = WSWelcome{}
+
+		var msg []byte
+		err := websocket.Message.Receive(ws, &msg)
+		if err != nil {
+			log.Println(err)
+			return
+		}
+
+		err = json.Unmarshal(msg, &welcome)
+		if err == nil {
+			WelcomeAgents(ws, welcome, "rr")
+		}
+
 		var lastip string
 		for {
 			var apievent = models.APIEvent{}
 			var event = models.Event{}
 
-			// Read
 			var msg []byte
 			err := websocket.Message.Receive(ws, &msg)
+			log.Println(string(msg))
 			if err != nil {
 				log.Println(err)
 				return
 			}
+			websocket.Message.Send(ws, "OK")
+
 			err = json.Unmarshal(msg, &apievent)
 			if err != nil {
 				log.Println(err)
@@ -45,8 +53,14 @@ func HandleWS(c echo.Context, cfg *config.Config) error {
 
 			switch apievent.MsgType {
 			case "bootstrap":
-				LISTENERS[apievent.Hostname] = ws
+				for _, value := range LISTENERS {
-				fmt.Println(LISTENERS)
+					err = websocket.Message.Send(value.ConnectionPS, msg)
+					if err != nil {
+						log.Println(err)
+					}
+				}
+
+				log.Println(LISTENERS)
 			case "add":
 				session := cfg.Db.NewSession()
 				event.APIParse(session, apievent)
@@ -56,16 +70,26 @@ func HandleWS(c echo.Context, cfg *config.Config) error {
 				if err != nil {
 					log.Println(err)
 				}
-				log.Printf("zmq: Inserted event")
+
+				for _, value := range LISTENERS {
+					err = websocket.Message.Send(value.ConnectionPS, msg)
+					if err != nil {
+						log.Println(err)
+					}
+				}
+				log.Printf("websocket: Inserted event")
 			case "file":
 				apievent.MsgType = "ws"
+			default:
+				log.Println("test")
 			}
 
-			// Write
+			WebSocketChannel <- apievent
-			err = websocket.Message.Send(ws, "OK_sendtest")
+
+			err = websocket.Message.Send(ws, "OK")
 			if err != nil {
-				gcconn(ws)
+				gcConn(ws, "rr")
-				fmt.Println(LISTENERS)
+				log.Println(LISTENERS)
 				return
 			}
 
@@ -73,12 +97,3 @@ func HandleWS(c echo.Context, cfg *config.Config) error {
 	}).ServeHTTP(c.Response(), c.Request())
 	return nil
 }
-
-func gcconn(conn *websocket.Conn) (err error) {
-	for i, v := range LISTENERS {
-		if v == conn {
-			delete(LISTENERS, i)
-		}
-	}
-	return err
-}