diff --git a/agents/agent.go b/agents/agent.go
index c9b919c..99f815b 100644
--- a/agents/agent.go
+++ b/agents/agent.go
@@ -37,10 +37,9 @@ type Message struct {
Payload string `json:"payload"`
}
-
var conn *websocket.Conn
-func registerAgent(agentName string, agentId string, agentIp string) error {
+func registerAgent(agentName, agentId, agentIp, agentType string) error {
// agent:= Agent{
// AgentName: agentName,
// InitialContact: time.Now().Format(time.RFC3339),
@@ -57,6 +56,7 @@ func registerAgent(agentName string, agentId string, agentIp string) error {
form := url.Values{}
form.Add("agentId", agentId)
form.Add("agentName", agentName)
+ form.Add("agentType", agentType)
form.Add("IPv4Address", agentIp)
resp, err := http.PostForm(registerURL, form)
@@ -124,13 +124,6 @@ func listenForCommands(agentName, agentIp string) {
continue
}
- // for {
- // _, rawMessage, err := conn.ReadMessage()
- // if err != nil {
- // log.Printf("Error reading message: %v", err)
- // break
- // }
-
var message Message
if err := json.Unmarshal(rawMessage, &message); err != nil {
log.Printf("Error unmarshalling message: %v", err)
@@ -180,8 +173,9 @@ func main() {
// agentId := "1234"
agentId := strconv.Itoa(randomInt(5))
agentIp := "127.0.0.1"
+ agentType := "BaseAgent"
- if err := registerAgent(agentName, agentId, agentIp); err != nil {
+ if err := registerAgent(agentName, agentId, agentIp, agentType); err != nil {
log.Fatalf("Agent registration failed: %v", err)
}
diff --git a/gomatic.sql b/gomatic.sql
index 82da9fa..fb896a0 100644
--- a/gomatic.sql
+++ b/gomatic.sql
@@ -4,6 +4,7 @@ drop table if exists agents;
create table agents (
id UUID default uuid() Primary Key,
agentId int unique,
+ agentType varchar(255),
agentName varchar(255),
IPv4Address varchar(15),
initialContact timestamp,
diff --git a/main.go b/main.go
index a90fe1c..1e42456 100644
--- a/main.go
+++ b/main.go
@@ -135,7 +135,6 @@ func getHomepage(w http.ResponseWriter, r *http.Request) {
}
func listAgents(w http.ResponseWriter, r *http.Request) {
- // agents, err := getAgents()
agents, err := api.GetAgents(db)
if err != nil {
http.Error(w, "Failed to fetch agents", http.StatusInternalServerError)
@@ -161,10 +160,6 @@ func main() {
var wg sync.WaitGroup
- // webSocketHandler := webSocketHandler {
- // upgrader: websocket.Upgrader{},
- // }
-
websocketServer := websocketserver.Server()
@@ -185,10 +180,6 @@ func main() {
webServer := &http.Server {
Addr: ":3333",
Handler: webMux,
- // BaseContext: func(l net.Listener) context.Context {
- // ctx = context.WithValue(ctx, keyServerAddr, l.Addr().String())
- // return ctx
- // },
}
wg.Add(1)
diff --git a/src/server/api/agentApi.go b/src/server/api/agentApi.go
index 7b2befc..4dbcf31 100644
--- a/src/server/api/agentApi.go
+++ b/src/server/api/agentApi.go
@@ -11,6 +11,7 @@ import (
type Agent struct {
AgentID int `json:"agentId"`
AgentName string `json:"agentName"`
+ AgentType string `json:"agentType"`
InitialContact string `json:"initialContact"`
LastContact string `json:"lastContact"`
IPv4Address string `json:"IPv4Address"`
@@ -44,12 +45,13 @@ func CreateAgent(db *sql.DB, w http.ResponseWriter, r * http.Request) (http.Resp
agentName := r.FormValue("agentName")
agentId := r.FormValue("agentId")
+ agentType := r.FormValue("agentType")
IPv4Address := r.FormValue("IPv4Address")
// initalContact := r.FormValue("initialContact")
// lastContact := r.FormValue("lastContact")
- query := "INSERT INTO agents (agentId, agentName, IPv4Address, initialContact, lastContact) VALUES (?, ?, ?, NOW(), NOW())"
- _, err = db.Exec(query, agentId, agentName, IPv4Address)
+ query := "INSERT INTO agents (agentId, agentName, agentType, IPv4Address, initialContact, lastContact) VALUES (?, ?, ?, ?, NOW(), NOW())"
+ _, err = db.Exec(query, agentId, agentName, agentType, IPv4Address)
if err != nil {
http.Error(w, "Failed to create agent", http.StatusInternalServerError)
return nil, err
@@ -81,7 +83,7 @@ func UpdateAgent(db *sql.DB, w http.ResponseWriter, r *http.Request, agentId str
}
func GetAgents(db *sql.DB) ([]Agent, error) {
- query := "SELECT agentId, agentName, IPv4Address, initialContact, lastContact FROM agents"
+ query := "SELECT agentId, agentName, agentType, IPv4Address, initialContact, lastContact FROM agents"
rows, err := db.Query(query)
if err != nil {
return nil, err
@@ -91,7 +93,7 @@ func GetAgents(db *sql.DB) ([]Agent, error) {
var agents []Agent
for rows.Next() {
var agent Agent
- err := rows.Scan(&agent.AgentID, &agent.AgentName, &agent.IPv4Address, &agent.InitialContact, &agent.LastContact)
+ err := rows.Scan(&agent.AgentID, &agent.AgentName, &agent.AgentType, &agent.IPv4Address, &agent.InitialContact, &agent.LastContact)
if err != nil {
return nil, err
}
@@ -101,9 +103,9 @@ func GetAgents(db *sql.DB) ([]Agent, error) {
}
func GetAgent(db *sql.DB, w http.ResponseWriter, r *http.Request, agentId string) (Agent, error) {
- query := "Select agentId, agentName, initialContact, lastContact from agents where agentId = ?"
+ query := "Select agentId, agentName, agentType, initialContact, lastContact from agents where agentId = ?"
var agent Agent
- err := db.QueryRow(query, agentId).Scan(&agent.AgentID, &agent.AgentName, &agent.InitialContact, &agent.LastContact)
+ err := db.QueryRow(query, agentId).Scan(&agent.AgentID, &agent.AgentName, &agent.AgentType, &agent.InitialContact, &agent.LastContact)
if err == sql.ErrNoRows {
http.Error(w, "Agent not found", http.StatusNotFound)
return Agent{} , err
diff --git a/src/server/websocket/websocketServer.go b/src/server/websocket/websocketServer.go
index 0399620..807f5b5 100644
--- a/src/server/websocket/websocketServer.go
+++ b/src/server/websocket/websocketServer.go
@@ -6,9 +6,13 @@ import (
"net/http"
"sync"
+ "time"
+
"github.com/gorilla/websocket"
)
+var responseChannels sync.Map // Key: agentName, Value: chan string
+
type webSocketHandler struct {
upgrader websocket.Upgrader
}
@@ -29,46 +33,6 @@ var getAgentNames http.HandlerFunc = func(w http.ResponseWriter, r *http.Request
}
-// func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-// c, err := wsh.upgrader.Upgrade(w, r, nil)
-// if err != nil {
-// log.Printf("Error %s when upgrading connection to websocket", err)
-// return
-// }
-// defer c.Close()
-
-// _, agentNameBytes, err := c.ReadMessage()
-// if err != nil {
-// log.Printf("Failed to read agent name: %s", err)
-// return
-// }
-
-// agentName := string(agentNameBytes)
-// agentSocketsMutex.Lock()
-// agentSockets[agentName] = c
-// agentSocketsMutex.Unlock()
-
-// log.Printf("Agent registered: %s", agentName)
-
-// for {
-// mt , message, err := c.ReadMessage()
-// if err != nil {
-// log.Printf("Error reading message: %s from agent: %s", err, agentName)
-// }
-
-// log.Printf("Received message: %s from agent: %s", message, agentName)
-// if err = c.WriteMessage(mt, message); err !=nil {
-// log.Printf("Error writing the message: %s", err)
-// break
-// }
-// }
-
-// agentSocketsMutex.Lock()
-// delete(agentSockets, agentName)
-// agentSocketsMutex.UnLock()
-// log.Printf("Agent disconnected: %s", agentName)
-// }
-
func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){
c, err := wsh.upgrader.Upgrade(w, r, nil)
if err != nil {
@@ -99,16 +63,17 @@ func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){
}()
for {
- mt, message, err := c.ReadMessage()
+ _, message, err := c.ReadMessage()
if err != nil {
log.Printf("Error reading from agent %s: %v", agentName, err)
break
}
log.Printf("Message from agent %s: %s", agentName, message)
- if err = c.WriteMessage(mt, message); err != nil {
- log.Printf("Error writing to agent %s: %v", agentName, err)
- break
+
+ if ch, ok := responseChannels.Load(agentName); ok {
+ responseChan := ch.(chan string)
+ responseChan <- string(message)
}
}
}
@@ -137,6 +102,10 @@ var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Reques
return
}
+ responseChan := make(chan string, 1)
+ responseChannels.Store(agentName, responseChan)
+ defer responseChannels.Delete(agentName)
+
message := Message {
Type: "command",
Payload: command,
@@ -150,8 +119,24 @@ var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Reques
return
}
- w.WriteHeader(http.StatusOK)
- w.Write([]byte("Command sent successfully"))
+ select {
+ case response := <-responseChan:
+ var parsedResponse map[string]string
+ if err := json.Unmarshal([]byte(response), &parsedResponse); err != nil {
+ http.Error(w, "Failed to parse response", http.StatusInternalServerError)
+ return
+ }
+ payload, ok := parsedResponse["payload"]
+ if !ok {
+ http.Error(w, "Invalid response structure", http.StatusInternalServerError)
+ return
+ }
+ w.WriteHeader(http.StatusOK)
+ w.Header().Set("Content-Type", "text/plain")
+ w.Write([]byte(payload))
+ case <- time.After(10 * time.Second):
+ http.Error(w, "Agent repsonse timed out", http.StatusGatewayTimeout)
+ }
}
func Server() (*http.Server) {
diff --git a/templates/index.html b/templates/index.html
index 649f620..b71078d 100644
--- a/templates/index.html
+++ b/templates/index.html
@@ -1,3 +1,288 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
@@ -37,7 +322,9 @@
const message = JSON.parse(event.data);
if (message.type === 'response') {
const output = document.getElementById('commandOutput');
- output.textContent = message.payload;
+ output.textContent = "";
+ output.innerText = message.payload.trim();
+ console.log("Raw websocket Data:", event.data);
}
};
diff --git a/templates/partials/agent_detail.html b/templates/partials/agent_detail.html
index f9f0f0c..7ac127d 100644
--- a/templates/partials/agent_detail.html
+++ b/templates/partials/agent_detail.html
@@ -2,6 +2,7 @@
Agent Details
ID: {{.AgentID}}
Name: {{.AgentName}}
+ Type: {{.AgentType}}
Initial Contact: {{.InitialContact}}
Last Contact: {{.LastContact}}
diff --git a/templates/partials/agent_list.html b/templates/partials/agent_list.html
index a31202a..f065ffc 100644
--- a/templates/partials/agent_list.html
+++ b/templates/partials/agent_list.html
@@ -3,6 +3,7 @@
| ID |
Name |
+ Type |
IPv4 Address |
Initial Contact |
Last Contact |
@@ -14,6 +15,7 @@
| {{.AgentID}} |
{{.AgentName}} |
+ {{.AgentType}} |
{{.IPv4Address}} |
{{.InitialContact}} |
{{.LastContact}} |