updated ipbl websocket feat

This commit is contained in:
Paul 2023-04-09 01:41:52 +02:00
parent 0354fe3255
commit f5ed3ee455
5 changed files with 163 additions and 30 deletions

View File

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

View File

@ -155,8 +155,12 @@ func RegisterRoutes(e *echo.Echo, ctx *context.Context, cfg *config.Config) {
return Result(c, err, disc) return Result(c, err, disc)
}) })
e.File("/test.html", "/home/paul/test.html") 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 return
}) })
} }

63
src/ws/init.go Normal file
View File

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

51
src/ws/pubsub.go Normal file
View File

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

View File

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