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

31
src/agentconnector/agentconnector.go
View File

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