used random names to identify the agents for execution of commands, moved agentCreate step from the agent itself to the websocket server. THis means a more simple agent is possible, because it only needs to do a single request to the websocket. Cleanup needed now
This commit is contained in:
Stefan Friese
parent
54604ff488
commit
c0764ac41d
|
|
@ -61,8 +61,8 @@ func registerAgent(agentName, agentId, agentIp, agentType string) error {
|
|||
}
|
||||
|
||||
|
||||
func connectToWebSocket(agentName, agentIp string) error {
|
||||
wsURL := fmt.Sprintf("ws://%s/data?agentName=%s&IPv4Address=%s", webSocketAddr, url.QueryEscape(agentName), url.QueryEscape(agentIp))
|
||||
func connectToWebSocket(agentName, agentId, agentIp, agentType string) error {
|
||||
wsURL := fmt.Sprintf("ws://%s/data?agentName=%s&agentId=%s&IPv4Address=%s&agentType=%s", webSocketAddr, url.QueryEscape(agentName), url.QueryEscape(agentId), url.QueryEscape(agentIp), url.QueryEscape(agentType))
|
||||
var err error
|
||||
for {
|
||||
conn, _, err = websocket.DefaultDialer.Dial(wsURL, nil)
|
||||
|
|
@ -76,13 +76,13 @@ func connectToWebSocket(agentName, agentIp string) error {
|
|||
}
|
||||
}
|
||||
|
||||
func reconnectToWebSocket(agentName, agentIp string) error {
|
||||
func reconnectToWebSocket(agentName, agentId, agentIp, agentType string) error {
|
||||
backoff := 2 * time.Second
|
||||
maxBackoff := 1 * time.Minute
|
||||
|
||||
for {
|
||||
log.Println("Attempting to reconnect to WebSocket...")
|
||||
err := connectToWebSocket(agentName, agentIp)
|
||||
err := connectToWebSocket(agentName, agentId, agentIp, agentType)
|
||||
if err == nil {
|
||||
log.Println("Reconnection succesful.")
|
||||
return nil
|
||||
|
|
@ -98,14 +98,14 @@ func reconnectToWebSocket(agentName, agentIp string) error {
|
|||
}
|
||||
}
|
||||
|
||||
func listenForCommands(agentName, agentIp string) {
|
||||
func listenForCommands(agentName, agentId, agentIp, agentType string) {
|
||||
defer conn.Close()
|
||||
|
||||
for {
|
||||
_, rawMessage, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
log.Printf("Connection lost: %v", err)
|
||||
if reconnectErr := reconnectToWebSocket(agentName, agentIp); reconnectErr != nil {
|
||||
if reconnectErr := reconnectToWebSocket(agentName, agentId, agentIp, agentType); reconnectErr != nil {
|
||||
log.Printf("Critical error during reconnection: %v", reconnectErr)
|
||||
}
|
||||
continue
|
||||
|
|
@ -161,13 +161,15 @@ func main() {
|
|||
agentIp := "127.0.0.1"
|
||||
agentType := "BaseAgent"
|
||||
|
||||
if err := registerAgent(agentName, agentId, agentIp, agentType); err != nil {
|
||||
log.Fatalf("Agent registration failed: %v", err)
|
||||
}
|
||||
log.Printf("AgentId: %s", agentId)
|
||||
|
||||
if err := connectToWebSocket(agentName, agentIp); err != nil {
|
||||
// if err := registerAgent(agentName, agentId, agentIp, agentType); err != nil {
|
||||
// log.Fatalf("Agent registration failed: %v", err)
|
||||
// }
|
||||
|
||||
if err := connectToWebSocket(agentName, agentId, agentIp, agentType); err != nil {
|
||||
log.Fatalf("Websocket connection failed: %v", err)
|
||||
}
|
||||
|
||||
listenForCommands(agentName, agentIp)
|
||||
listenForCommands(agentName, agentId, agentIp, agentType)
|
||||
}
|
||||
|
|
|
|||
7
go.mod
7
go.mod
|
|
@ -3,9 +3,14 @@ module gontrol
|
|||
go 1.23.4
|
||||
|
||||
require (
|
||||
github.com/PuerkitoBio/goquery v1.10.1
|
||||
github.com/go-sql-driver/mysql v1.8.1
|
||||
github.com/gorilla/websocket v1.5.3
|
||||
github.com/kelseyhightower/envconfig v1.4.0
|
||||
)
|
||||
|
||||
require filippo.io/edwards25519 v1.1.0 // indirect
|
||||
require (
|
||||
filippo.io/edwards25519 v1.1.0 // indirect
|
||||
github.com/andybalholm/cascadia v1.3.3 // indirect
|
||||
golang.org/x/net v0.33.0 // indirect
|
||||
)
|
||||
|
|
|
|||
70
go.sum
70
go.sum
|
|
@ -1,8 +1,78 @@
|
|||
filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
|
||||
filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=
|
||||
github.com/PuerkitoBio/goquery v1.10.1 h1:Y8JGYUkXWTGRB6Ars3+j3kN0xg1YqqlwvdTV8WTFQcU=
|
||||
github.com/PuerkitoBio/goquery v1.10.1/go.mod h1:IYiHrOMps66ag56LEH7QYDDupKXyo5A8qrjIx3ZtujY=
|
||||
github.com/andybalholm/cascadia v1.3.3 h1:AG2YHrzJIm4BZ19iwJ/DAua6Btl3IwJX+VI4kktS1LM=
|
||||
github.com/andybalholm/cascadia v1.3.3/go.mod h1:xNd9bqTn98Ln4DwST8/nG+H0yuB8Hmgu1YHNnWw0GeA=
|
||||
github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpvNJ1Y=
|
||||
github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg=
|
||||
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
|
||||
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
|
||||
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
|
||||
github.com/kelseyhightower/envconfig v1.4.0 h1:Im6hONhd3pLkfDFsbRgu68RDNkGF1r3dvMUtDTo2cv8=
|
||||
github.com/kelseyhightower/envconfig v1.4.0/go.mod h1:cccZRl6mQpaq41TPp5QxidR+Sa3axMbJDNb//FQX6Gg=
|
||||
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
|
||||
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
|
||||
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
|
||||
golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliYc=
|
||||
golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
|
||||
golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
|
||||
golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
|
||||
golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
|
||||
golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
|
||||
golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
|
||||
golang.org/x/mod v0.15.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
|
||||
golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
|
||||
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
|
||||
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
|
||||
golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
|
||||
golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
|
||||
golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
|
||||
golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk=
|
||||
golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
|
||||
golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
|
||||
golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I=
|
||||
golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
|
||||
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
|
||||
golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2/go.mod h1:TeRTkGYfJXctD9OcfyVLyj2J3IxLnKwHJR8f4D8a3YE=
|
||||
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
||||
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
|
||||
golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo=
|
||||
golang.org/x/term v0.12.0/go.mod h1:owVbMEjm3cBLCHdkQu9b1opXd4ETQWc3BhuQGKgXgvU=
|
||||
golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk=
|
||||
golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
|
||||
golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
|
||||
golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
|
||||
golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
|
||||
golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
|
||||
golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
|
||||
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
|
||||
golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
|
||||
golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
|
||||
golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58=
|
||||
golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
|
||||
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
|
|
|
|||
6
main.go
6
main.go
|
|
@ -129,6 +129,11 @@ func getAgentNames(w http.ResponseWriter, r *http.Request) {
|
|||
return
|
||||
}
|
||||
|
||||
func getAgentIds(w http.ResponseWriter, r *http.Request) {
|
||||
api.GetAgentIds(db, w, r)
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
func main() {
|
||||
|
||||
|
|
@ -146,6 +151,7 @@ func main() {
|
|||
webMux.HandleFunc("/", getHomepage)
|
||||
webMux.HandleFunc("/agents", agentsHandler)
|
||||
webMux.HandleFunc("/agentNames", getAgentNames)
|
||||
webMux.HandleFunc("/agentIds", getAgentIds)
|
||||
webMux.HandleFunc("/agents/{agentId}", agentsHandler)
|
||||
|
||||
initDB (cfg.Database.Username, cfg.Database.Password, cfg.Database.Host, cfg.Database.Port, cfg.Database.Name)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ import (
|
|||
"net/http"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"strconv"
|
||||
"log"
|
||||
|
||||
_ "github.com/go-sql-driver/mysql"
|
||||
)
|
||||
|
|
@ -53,6 +55,7 @@ func CreateAgent(db *sql.DB, w http.ResponseWriter, r * http.Request) (http.Resp
|
|||
query := "INSERT INTO agents (agentId, agentName, agentType, IPv4Address, initialContact, lastContact) VALUES (?, ?, ?, ?, NOW(), NOW())"
|
||||
_, err = db.Exec(query, agentId, agentName, agentType, IPv4Address)
|
||||
if err != nil {
|
||||
log.Printf("Database err is: %s", err)
|
||||
http.Error(w, "Failed to create agent", http.StatusInternalServerError)
|
||||
return nil, err
|
||||
}
|
||||
|
|
@ -140,3 +143,27 @@ func GetAgentNames(db *sql.DB, w http.ResponseWriter, r *http.Request) {
|
|||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(agentNames)
|
||||
}
|
||||
|
||||
|
||||
func GetAgentIds(db *sql.DB, w http.ResponseWriter, r *http.Request) {
|
||||
query := "SELECT agentId from agents"
|
||||
rows, err := db.Query(query)
|
||||
if err != nil {
|
||||
http.Error(w, "Failed to fetch agent names", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var agentIds []string
|
||||
for rows.Next() {
|
||||
var id int
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
http.Error(w, "Error reading agent names", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
agentIds = append(agentIds, strconv.Itoa(id))
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(agentIds)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,16 +1,26 @@
|
|||
package websocketserver
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"gontrol/src/randomname"
|
||||
"gontrol/src/server/api"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/PuerkitoBio/goquery"
|
||||
_ "github.com/go-sql-driver/mysql"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
var responseChannels sync.Map // Key: agentName, Value: chan string
|
||||
var db *sql.DB
|
||||
|
||||
type webSocketHandler struct {
|
||||
upgrader websocket.Upgrader
|
||||
|
|
@ -32,6 +42,111 @@ var getAgentNames http.HandlerFunc = func(w http.ResponseWriter, r *http.Request
|
|||
|
||||
}
|
||||
|
||||
func registerAgent(agentName, agentId, agentIp, agentType string) error {
|
||||
|
||||
registerURL := "http://localhost:3333/agents"
|
||||
|
||||
form := url.Values{}
|
||||
form.Add("agentId", agentId)
|
||||
form.Add("agentName", agentName)
|
||||
form.Add("agentType", agentType)
|
||||
form.Add("IPv4Address", agentIp)
|
||||
|
||||
resp, err := http.PostForm(registerURL, form)
|
||||
if err != nil {
|
||||
log.Printf("Error registering agent: %v", err)
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode == http.StatusCreated {
|
||||
log.Printf("Agent %s successfully registered.", agentName)
|
||||
return nil
|
||||
} else if resp.StatusCode == http.StatusOK {
|
||||
log.Printf("Agent %s already registered.", agentName)
|
||||
return nil
|
||||
} else {
|
||||
log.Printf("Failed to register agent, status: %v", resp.Status)
|
||||
return err
|
||||
}
|
||||
|
||||
// log.Printf("Agent %s successfully registered.", agentName)
|
||||
}
|
||||
|
||||
func getAgentDetails(agentId string) (*api.Agent, error) {
|
||||
agentURL := "http://localhost:3333/agents/" + agentId
|
||||
// var ids []string
|
||||
resp, err := http.Get(agentURL)
|
||||
if err != nil {
|
||||
log.Printf("Failed to make GET request: %w", err)
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
doc, err := goquery.NewDocumentFromReader(resp.Body)
|
||||
if err != nil {
|
||||
log.Printf("Failed to parse HTML: %w", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
agent := &api.Agent{}
|
||||
doc.Find("#agent-detail p").Each(func(i int, s *goquery.Selection) {
|
||||
text := s.Text()
|
||||
if strings.HasPrefix(text, "ID:") {
|
||||
agent.AgentID, err = strconv.Atoi(strings.TrimSpace(strings.TrimPrefix(text, "ID:")))
|
||||
if err != nil {
|
||||
log.Printf("Converting string to integer failed in getAgentDetails(): %w", err)
|
||||
}
|
||||
} else if strings.HasPrefix(text, "Name:") {
|
||||
agent.AgentName = strings.TrimSpace(strings.TrimPrefix(text, "Name:"))
|
||||
} else if strings.HasPrefix(text, "Type:") {
|
||||
agent.AgentType = strings.TrimSpace(strings.TrimPrefix(text, "Type:"))
|
||||
}
|
||||
})
|
||||
return agent, nil
|
||||
}
|
||||
|
||||
func containsId(ids []string, agentId string) bool {
|
||||
for _, id := range ids {
|
||||
if id == agentId {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func getAgentIds() ([]string, error) {
|
||||
|
||||
idURL := "http://localhost:3333/agentIds"
|
||||
// var ids []string
|
||||
resp, err := http.Get(idURL)
|
||||
if err != nil {
|
||||
log.Printf("Failed to make GET request: %w", err)
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Printf("Unexpected status code: %d", resp.StatusCode)
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
log.Printf("Failed to read response body: %w", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var agentIds []string
|
||||
if err := json.Unmarshal(body, &agentIds); err != nil {
|
||||
log.Printf("Failed to parse JSON response: %w", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return agentIds, nil
|
||||
}
|
||||
|
||||
|
||||
func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){
|
||||
c, err := wsh.upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
|
|
@ -39,15 +154,32 @@ func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){
|
|||
return
|
||||
}
|
||||
|
||||
agentName := r.URL.Query().Get("agentName")
|
||||
agentIP := r.URL.Query().Get("IPv4Address")
|
||||
agentId := r.URL.Query().Get("agentId")
|
||||
agentType := r.URL.Query().Get("agentType")
|
||||
agentName := ""
|
||||
|
||||
agentIds, err := getAgentIds()
|
||||
if err != nil {
|
||||
log.Printf("Error %v\n", err)
|
||||
return
|
||||
}
|
||||
|
||||
if !containsId(agentIds, agentId) {
|
||||
agentName = randomname.GenerateRandomName()
|
||||
registerAgent(agentName, agentId, agentIP, agentType)
|
||||
} else {
|
||||
agentDetails, _ := getAgentDetails(agentId)
|
||||
agentName = agentDetails.AgentName
|
||||
}
|
||||
|
||||
if agentName == "" || agentIP == "" {
|
||||
log.Printf("Missing agentName or IPv4Address in query parameters")
|
||||
c.Close()
|
||||
return
|
||||
}
|
||||
|
||||
log.Printf("Agent connected: %s (%s)", agentName, agentIP)
|
||||
log.Printf("Agent %s connected: %s (%s)", agentId, agentName, agentIP)
|
||||
|
||||
agentSocketsMutex.Lock()
|
||||
agentSockets[agentName] = c
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@
|
|||
})
|
||||
.catch(error => console.error('Error fetching agent names:', error));
|
||||
|
||||
const socket = new WebSocket("ws://localhost:5555/data");
|
||||
const socket = new WebSocket("ws://localhost:5555/executeCommand");
|
||||
socket.onmessage = (event) => {
|
||||
const message = JSON.parse(event.data);
|
||||
if (message.type === 'response') {
|
||||
|
|
|
|||
Loading…
Reference in New Issue