whx
/
gommand
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

added hostname parameter

This commit is contained in:
Stefan Etringer 2025-05-26 15:06:49 +00:00
parent 380ff26eef
commit 20afb2f216
2 changed files with 24 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
main.go
View File

@ -491,7 +491,7 @@ func handler(w http.ResponseWriter, r *http.Request) {
tmpl.Execute(w, data) tmpl.Execute(w, data)
} }
func startMainServer() (int, net.Listener) { func startInteractiveServer() (int, net.Listener) {
http.HandleFunc("/", handler) http.HandleFunc("/", handler)
http.HandleFunc("/upload", fileUploadHandler) http.HandleFunc("/upload", fileUploadHandler)
@ -529,7 +529,7 @@ func main() {
// log.Println("Using port:", listener.Addr().(*net.TCPAddr).Port) // log.Println("Using port:", listener.Addr().(*net.TCPAddr).Port)
// log.Fatal(http.Serve(listener, nil)) // log.Fatal(http.Serve(listener, nil))
port, listener := startMainServer() port, listener := startInteractiveServer()
var wg sync.WaitGroup var wg sync.WaitGroup
wg.Add(2) wg.Add(2)

31
src/agentconnector/agentconnector.go
View File

@ -7,6 +7,7 @@ import (
"net/http" "net/http"
"net/url" "net/url"
"os/exec" "os/exec"
"os"
"time" "time"
"math/rand" "math/rand"
"math" "math"
@ -30,6 +31,7 @@ type Agent struct {
InitialContact string `json:"initialContact"` InitialContact string `json:"initialContact"`
LastContact string `json:"lastContact"` LastContact string `json:"lastContact"`
AddPort string `json:"addPort"` AddPort string `json:"addPort"`
HostName string `json:"HostName"`
} }
type Message struct { type Message struct {
@ -39,7 +41,7 @@ type Message struct {
var conn *websocket.Conn var conn *websocket.Conn
func registerAgent(agentName, agentId, agentIp, agentType, addPort string) error { func registerAgent(agentName, agentId, agentIp, agentType, addPort, hostname string) error {
form := url.Values{} form := url.Values{}
form.Add("agentId", agentId) form.Add("agentId", agentId)
@ -47,6 +49,7 @@ func registerAgent(agentName, agentId, agentIp, agentType, addPort string) error
form.Add("agentType", agentType) form.Add("agentType", agentType)
form.Add("IPv4Address", agentIp) form.Add("IPv4Address", agentIp)
form.Add("addPort", addPort) form.Add("addPort", addPort)
form.Add("hostname", hostname)
resp, err := http.PostForm(registerURL, form) resp, err := http.PostForm(registerURL, form)
if err != nil { if err != nil {
@ -62,8 +65,17 @@ func registerAgent(agentName, agentId, agentIp, agentType, addPort string) error
return nil return nil
} }
func connectToWebSocket(agentName, agentId, agentIp, agentType, addPort string) error { func connectToWebSocket(agentName, agentId, agentIp, agentType, addPort, hostname string) error {
wsURL := fmt.Sprintf("ws://%s/data?agentName=%s&agentId=%s&IPv4Address=%s&agentType=%s&addPort=%s", webSocketAddr, url.QueryEscape(agentName), url.QueryEscape(agentId), url.QueryEscape(agentIp), url.QueryEscape(agentType), url.QueryEscape(addPort)) wsURL := fmt.Sprintf(
"ws://%s/data?agentName=%s&agentId=%s&IPv4Address=%s&agentType=%s&addPort=%s&hostname=%s",
webSocketAddr,
url.QueryEscape(agentName),
url.QueryEscape(agentId),
url.QueryEscape(agentIp),
url.QueryEscape(agentType),
url.QueryEscape(addPort),
url.QueryEscape(hostname),
)
var err error var err error
for { for {
conn, _, err = websocket.DefaultDialer.Dial(wsURL, nil) conn, _, err = websocket.DefaultDialer.Dial(wsURL, nil)
@ -78,13 +90,13 @@ func connectToWebSocket(agentName, agentId, agentIp, agentType, addPort string)
} }
} }
func reconnectToWebSocket(agentName, agentId, agentIp, agentType, addPort string) error { func reconnectToWebSocket(agentName, agentId, agentIp, agentType, addPort, hostname string) error {
backoff := 2 * time.Second backoff := 2 * time.Second
maxBackoff := 1 * time.Minute maxBackoff := 1 * time.Minute
for { for {
log.Println("Attempting to reconnect to WebSocket...") log.Println("Attempting to reconnect to WebSocket...")
err := connectToWebSocket(agentName, agentId, agentIp, agentType, addPort) err := connectToWebSocket(agentName, agentId, agentIp, agentType, addPort, hostname)
if err == nil { if err == nil {
log.Println("Reconnection succesful.") log.Println("Reconnection succesful.")
return nil return nil
@ -100,14 +112,14 @@ func reconnectToWebSocket(agentName, agentId, agentIp, agentType, addPort string
} }
} }
func listenForCommands(agentName, agentId, agentIp, agentType, addPort string) { func listenForCommands(agentName, agentId, agentIp, agentType, addPort, hostname string) {
defer conn.Close() defer conn.Close()
for { for {
_, rawMessage, err := conn.ReadMessage() _, rawMessage, err := conn.ReadMessage()
if err != nil { if err != nil {
log.Printf("Connection lost: %v", err) log.Printf("Connection lost: %v", err)
if reconnectErr := reconnectToWebSocket(agentName, agentId, agentIp, agentType, addPort); reconnectErr != nil { if reconnectErr := reconnectToWebSocket(agentName, agentId, agentIp, agentType, addPort, hostname); reconnectErr != nil {
log.Printf("Critical error during reconnection: %v", reconnectErr) log.Printf("Critical error during reconnection: %v", reconnectErr)
} }
continue continue
@ -165,6 +177,7 @@ func StartServer(agentInteractivePort int){
agentIp := "127.0.0.1" agentIp := "127.0.0.1"
agentType := "Interactive" agentType := "Interactive"
addPort := strconv.Itoa(agentInteractivePort) addPort := strconv.Itoa(agentInteractivePort)
hostname, _ := os.Hostname()
log.Printf("AgentId: %s", agentId) log.Printf("AgentId: %s", agentId)
@ -174,9 +187,9 @@ func StartServer(agentInteractivePort int){
// log.Fatalf("Agent registration failed: %v", err) // log.Fatalf("Agent registration failed: %v", err)
// } // }
if err := connectToWebSocket(agentName, agentId, agentIp, agentType, addPort); err != nil { if err := connectToWebSocket(agentName, agentId, agentIp, agentType, addPort, hostname); err != nil {
log.Fatalf("Websocket connection failed: %v", err) log.Fatalf("Websocket connection failed: %v", err)
} }
listenForCommands(agentName, agentId, agentIp, agentType, addPort) listenForCommands(agentName, agentId, agentIp, agentType, addPort, hostname)
} }