From 697df686404751aa7e327567761b5791dcec0a33 Mon Sep 17 00:00:00 2001
From: Paul Lecuq <paul@paulbsd.com>
Date: Sun, 5 Mar 2023 23:05:05 +0100
Subject: [PATCH] updated ws feature

---
 cmd/ipbl/ipbl.go     |  2 ++
 src/routers/funcs.go |  3 +--
 src/ws/main.go       | 52 ++++++++++++++++++++++++++++++++++++++++++++
 src/ws/pubsub.go     | 27 -----------------------
 src/ws/reqrep.go     | 31 --------------------------
 5 files changed, 55 insertions(+), 60 deletions(-)
 delete mode 100644 src/ws/pubsub.go
 delete mode 100644 src/ws/reqrep.go

diff --git a/cmd/ipbl/ipbl.go b/cmd/ipbl/ipbl.go
index 5d32e21..02fa911 100644
--- a/cmd/ipbl/ipbl.go
+++ b/cmd/ipbl/ipbl.go
@@ -10,6 +10,7 @@ import (
 	"git.paulbsd.com/paulbsd/ipbl/src/database"
 	"git.paulbsd.com/paulbsd/ipbl/src/models"
 	"git.paulbsd.com/paulbsd/ipbl/src/routers"
+	"git.paulbsd.com/paulbsd/ipbl/src/ws"
 	"git.paulbsd.com/paulbsd/ipbl/utils"
 	_ "github.com/lib/pq"
 )
@@ -36,6 +37,7 @@ func main() {
 		go models.ScanIP(&cfg)
 	}
 	//go zmqrouter.Init(&cfg)
+	go ws.Init(&cfg)
 	go func() { err = routers.RunServer(&ctx, &cfg) }()
 	if err != nil {
 		log.Fatalln(err)
diff --git a/src/routers/funcs.go b/src/routers/funcs.go
index 1d8e2b4..cabf1b9 100644
--- a/src/routers/funcs.go
+++ b/src/routers/funcs.go
@@ -155,6 +155,5 @@ 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("/wsps", ws.HandleWSPubSub)
-	e.GET("/wsrr", ws.HandleWSReqRep)
+	e.GET("/ws", ws.HandleWS)
 }
diff --git a/src/ws/main.go b/src/ws/main.go
index 9859295..12aea25 100644
--- a/src/ws/main.go
+++ b/src/ws/main.go
@@ -1 +1,53 @@
 package ws
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+
+	"git.paulbsd.com/paulbsd/ipbl/src/config"
+	"git.paulbsd.com/paulbsd/ipbl/src/models"
+	"github.com/labstack/echo/v4"
+	"golang.org/x/net/websocket"
+)
+
+var LISTENERS map[string]*websocket.Conn
+var CHAN int
+
+func Init(cfg *config.Config) {
+	LISTENERS = make(map[string]*websocket.Conn)
+}
+
+func HandleWS(c echo.Context) error {
+	websocket.Handler(func(ws *websocket.Conn) {
+		defer ws.Close()
+
+		fmt.Printf("rr: %v\n", ws)
+
+		for {
+			// Read
+			var event models.APIEvent
+			var msg []byte
+			err := websocket.Message.Receive(ws, &msg)
+			if err != nil {
+				return
+			}
+			err = json.Unmarshal(msg, &event)
+			if err != nil {
+				log.Println(err)
+			}
+			if event.MsgType == "bootstrap" {
+				LISTENERS[event.Hostname] = ws
+			}
+			fmt.Println(LISTENERS)
+
+			// Write
+			err = websocket.Message.Send(ws, "OK")
+			if err != nil {
+				return
+			}
+
+		}
+	}).ServeHTTP(c.Response(), c.Request())
+	return nil
+}
diff --git a/src/ws/pubsub.go b/src/ws/pubsub.go
deleted file mode 100644
index fdacd55..0000000
--- a/src/ws/pubsub.go
+++ /dev/null
@@ -1,27 +0,0 @@
-package ws
-
-import (
-	"github.com/labstack/echo/v4"
-	"golang.org/x/net/websocket"
-)
-
-func HandleWSPubSub(c echo.Context) error {
-	websocket.Handler(func(ws *websocket.Conn) {
-		defer ws.Close()
-		for {
-			// Write
-			err := websocket.Message.Send(ws, "Hello, Client!")
-			if err != nil {
-				return
-			}
-
-			// Read
-			msg := ""
-			err = websocket.Message.Receive(ws, &msg)
-			if err != nil {
-				return
-			}
-		}
-	}).ServeHTTP(c.Response(), c.Request())
-	return nil
-}
diff --git a/src/ws/reqrep.go b/src/ws/reqrep.go
deleted file mode 100644
index 89e0e27..0000000
--- a/src/ws/reqrep.go
+++ /dev/null
@@ -1,31 +0,0 @@
-package ws
-
-import (
-	"fmt"
-
-	"github.com/labstack/echo/v4"
-	"golang.org/x/net/websocket"
-)
-
-func HandleWSReqRep(c echo.Context) error {
-	websocket.Handler(func(ws *websocket.Conn) {
-		defer ws.Close()
-		for {
-			// Read
-			msg := ""
-			err := websocket.Message.Receive(ws, &msg)
-			if err != nil {
-				return
-			}
-			fmt.Println(msg)
-
-			// Write
-			err = websocket.Message.Send(ws, "OK")
-			if err != nil {
-				return
-			}
-
-		}
-	}).ServeHTTP(c.Response(), c.Request())
-	return nil
-}