solved issue with the style of presentation of the agent response on the website. Added agentType to the api, agent and database

This commit is contained in:
Stefan Friese 2025-01-22 16:02:44 +00:00
parent 323c04af89
commit 9e45adfff5
8 changed files with 335 additions and 72 deletions

View File

@ -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)
}

View File

@ -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,

View File

@ -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)

View File

@ -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

View File

@ -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
}
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.Write([]byte("Command sent successfully"))
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) {

View File

@ -1,3 +1,288 @@
<!-- <!DOCTYPE html> -->
<!-- <html lang="en"> -->
<!-- <head> -->
<!-- <meta charset="UTF-8"> -->
<!-- <meta http-equiv="X-UA-Compatible" content="IE=edge"> -->
<!-- <meta name="viewport" content="width=device-width, initial-scale=1.0"> -->
<!-- <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/css/bootstrap.min.css" rel="stylesheet" integrity="sha384-QWTKZyjpPEjISv5WaRU9OFeRpok6YctnYmDr5pNlyT2bRjXh0JMhjY6hW+ALEwIH" crossorigin="anonymous"> -->
<!-- <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/js/bootstrap.bundle.min.js"></script> -->
<!-- <script src="https://unpkg.com/htmx.org@1.9.12"></script> -->
<!-- <title>g2: gommand & gontrol</title> -->
<!-- <script> -->
<!-- document.addEventListener('DOMContentLoaded', () => { -->
<!-- // Fetch agent names to populate the dropdown -->
<!-- fetch('/agentNames') -->
<!-- .then(response => response.json()) -->
<!-- .then(agentNames => { -->
<!-- const dropdown = document.getElementById('agentName'); -->
<!-- agentNames.forEach(name => { -->
<!-- const option = document.createElement('option'); -->
<!-- option.value = name; -->
<!-- option.textContent = name; -->
<!-- dropdown.appendChild(option); -->
<!-- }); -->
<!-- }) -->
<!-- .catch(error => console.error('Error fetching agent names:', error)); -->
<!-- // WebSocket setup -->
<!-- const socket = new WebSocket("ws://localhost:5555/data"); -->
<!-- // Handle WebSocket messages -->
<!-- socket.onmessage = (event) => { -->
<!-- const message = JSON.parse(event.data); -->
<!-- if (message.type === 'response') { -->
<!-- const output = document.getElementById('commandOutput'); -->
<!-- output.textContent = message.payload; // Display the agent's response -->
<!-- } -->
<!-- }; -->
<!-- // Handle WebSocket errors -->
<!-- socket.onerror = (error) => { -->
<!-- console.error("WebSocket error:", error); -->
<!-- }; -->
<!-- // Handle WebSocket closure -->
<!-- socket.onclose = () => { -->
<!-- console.warn('WebSocket connection closed.'); -->
<!-- }; -->
<!-- // Intercept HTMX requests for /executeCommand -->
<!-- document.body.addEventListener('htmx:beforeRequest', (evt) => { -->
<!-- const url = evt.detail.pathInfo.requestPath; -->
<!-- if (url === '/executeCommand') { -->
<!-- evt.preventDefault(); // Stop HTMX from processing the request -->
<!-- // Collect form data -->
<!-- const formData = new FormData(evt.detail.elt); -->
<!-- const payload = new URLSearchParams(formData); // Convert form data to query string -->
<!-- // Send form data using fetch -->
<!-- fetch('http://localhost:5555/executeCommand', { -->
<!-- method: 'POST', -->
<!-- headers: { -->
<!-- 'Content-Type': 'application/x-www-form-urlencoded', -->
<!-- }, -->
<!-- body: payload, -->
<!-- }).then((response) => { -->
<!-- if (!response.ok) { -->
<!-- console.error('Failed to execute command:', response.statusText); -->
<!-- } -->
<!-- console.log(response) -->
<!-- }); -->
<!-- } -->
<!-- }); -->
<!-- }); -->
<!-- </script> -->
<!-- </head> -->
<!-- <body> -->
<!-- <div class="container"> -->
<!-- <div class="row"> -->
<!-- <div class="col"> -->
<!-- <h2>Agents</h2> -->
<!-- <!-1- Agent List -1-> -->
<!-- <div id="agentList" hx-get="/agents" hx-trigger="load" hx-swap="innerHTML"></div> -->
<!-- <!-1- Command Execution -1-> -->
<!-- <div id="agentCommands"> -->
<!-- <h3>Command Execution</h3> -->
<!-- <form hx-post="/executeCommand" hx-target="#commandOutput" hx-encoding="application/x-www-form-urlencoded" hx-swap="innerHTML"> -->
<!-- <div class="mb-3"> -->
<!-- <label for="agentName" class="form-label">Agent Name</label> -->
<!-- <select class="form-select" id="agentName" name="agentName" required> -->
<!-- <!-1- Dynamically populated with agent names -1-> -->
<!-- </select> -->
<!-- </div> -->
<!-- <div class="mb-3"> -->
<!-- <label for="command" class="form-label">Command</label> -->
<!-- <input type="text" class="form-control" id="command" name="command" placeholder="Enter command" required> -->
<!-- </div> -->
<!-- <button type="submit" class="btn btn-primary">Execute</button> -->
<!-- </form> -->
<!-- <div id="commandOutput" class="mt-3 p-2 border">Here goes the output...</div> -->
<!-- </div> -->
<!-- <!-1- Add Agent Form -1-> -->
<!-- <button class="btn btn-primary mt-3" data-bs-toggle="collapse" data-bs-target="#addAgentForm">Add Agent</button> -->
<!-- <div id="addAgentForm" class="collapse mt-2"> -->
<!-- <form hx-post="/agents" hx-target="#agentList" hx-swap="innerHTML"> -->
<!-- <div class="mb-3"> -->
<!-- <label for="agentId" class="form-label">Agent Id</label> -->
<!-- <input type="text" class="form-control" id="agentId" name="agentId" required> -->
<!-- </div> -->
<!-- <div class="mb-3"> -->
<!-- <label for="agentName" class="form-label">Agent Name</label> -->
<!-- <input type="text" class="form-control" id="agentNameInput" name="agentName" required> -->
<!-- </div> -->
<!-- <div class="mb-3"> -->
<!-- <label for="IPv4Address" class="form-label">IPv4 Address</label> -->
<!-- <input type="text" class="form-control" id="IPv4Address" name="IPv4Address" required> -->
<!-- </div> -->
<!-- <div class="mb-3"> -->
<!-- <label for="initialContact" class="form-label">Initial Contact</label> -->
<!-- <input type="datetime-local" class="form-control" id="initialContact" name="initialContact" required> -->
<!-- </div> -->
<!-- <div class="mb-3"> -->
<!-- <label for="lastContact" class="form-label">Last Contact</label> -->
<!-- <input type="datetime-local" class="form-control" id="lastContact" name="lastContact" required> -->
<!-- </div> -->
<!-- <button type="submit" class="btn btn-success">Add Agent</button> -->
<!-- </form> -->
<!-- </div> -->
<!-- </div> -->
<!-- <!-1- Agent Details -1-> -->
<!-- <div class="col" id="agentDetails"> -->
<!-- <h3>Details</h3> -->
<!-- <p>Select an agent to view details.</p> -->
<!-- </div> -->
<!-- </div> -->
<!-- </div> -->
<!-- </body> -->
<!-- </html> -->
<!-- <!DOCTYPE html> -->
<!-- <html lang="en"> -->
<!-- <head> -->
<!-- <meta charset="UTF-8"> -->
<!-- <meta http-equiv="X-UA-Compatible" content="IE=edge"> -->
<!-- <meta name="viewport" content="width=device-width, initial-scale=1.0"> -->
<!-- <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/css/bootstrap.min.css" rel="stylesheet" integrity="sha384-QWTKZyjpPEjISv5WaRU9OFeRpok6YctnYmDr5pNlyT2bRjXh0JMhjY6hW+ALEwIH" crossorigin="anonymous"> -->
<!-- <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/js/bootstrap.bundle.min.js"></script> -->
<!-- <script src="https://unpkg.com/htmx.org@1.9.12"></script> -->
<!-- <title>g2: gommand & gontrol</title> -->
<!-- <script> -->
<!-- document.addEventListener('DOMContentLoaded', () => { -->
<!-- // Fetch agent names to populate the dropdown -->
<!-- fetch('/agentNames') -->
<!-- .then(response => response.json()) -->
<!-- .then(agentNames => { -->
<!-- const dropdown = document.getElementById('agentName'); -->
<!-- agentNames.forEach(name => { -->
<!-- const option = document.createElement('option'); -->
<!-- option.value = name; -->
<!-- option.textContent = name; -->
<!-- dropdown.appendChild(option); -->
<!-- }); -->
<!-- }) -->
<!-- .catch(error => console.error('Error fetching agent names:', error)); -->
<!-- // WebSocket setup for listening to agent responses -->
<!-- const socket = new WebSocket("ws://localhost:5555/executeCommand"); -->
<!-- // Handle WebSocket messages -->
<!-- socket.onmessage = (event) => { -->
<!-- const message = JSON.parse(event.data); -->
<!-- if (message.type === 'response') { -->
<!-- const output = document.getElementById('commandOutput'); -->
<!-- output.textContent = message.payload; // Display the agent's response -->
<!-- } -->
<!-- }; -->
<!-- // Handle WebSocket errors -->
<!-- socket.onerror = (error) => { -->
<!-- console.error("WebSocket error:", error); -->
<!-- }; -->
<!-- // Handle WebSocket closure -->
<!-- socket.onclose = () => { -->
<!-- console.warn('WebSocket connection closed.'); -->
<!-- }; -->
<!-- // Intercept HTMX requests for /executeCommand to send via WebSocket -->
<!-- document.body.addEventListener('htmx:beforeRequest', (evt) => { -->
<!-- const url = evt.detail.pathInfo.requestPath; -->
<!-- if (url === '/executeCommand') { -->
<!-- evt.preventDefault(); // Stop HTMX from processing the request -->
<!-- // Collect form data -->
<!-- const formData = new FormData(evt.detail.elt); -->
<!-- const payload = new URLSearchParams(formData); // Convert form data to query string -->
<!-- // Send form data using fetch -->
<!-- fetch('http://localhost:5555/executeCommand', { -->
<!-- method: 'POST', -->
<!-- headers: { -->
<!-- 'Content-Type': 'application/x-www-form-urlencoded', -->
<!-- }, -->
<!-- body: payload, -->
<!-- }).then((response) => { -->
<!-- if (!response.ok) { -->
<!-- console.error('Failed to execute command:', response.statusText); -->
<!-- } -->
<!-- }); -->
<!-- } -->
<!-- }); -->
<!-- }); -->
<!-- </script> -->
<!-- </head> -->
<!-- <body> -->
<!-- <div class="container"> -->
<!-- <div class="row"> -->
<!-- <div class="col"> -->
<!-- <h2>Agents</h2> -->
<!-- <!-1- Agent List -1-> -->
<!-- <div id="agentList" hx-get="/agents" hx-trigger="load" hx-swap="innerHTML"></div> -->
<!-- <!-1- Command Execution -1-> -->
<!-- <div id="agentCommands"> -->
<!-- <h3>Command Execution</h3> -->
<!-- <form hx-post="/executeCommand" hx-target="#commandOutput" hx-encoding="application/x-www-form-urlencoded" hx-swap="innerHTML"> -->
<!-- <div class="mb-3"> -->
<!-- <label for="agentName" class="form-label">Agent Name</label> -->
<!-- <select class="form-select" id="agentName" name="agentName" required> -->
<!-- <!-1- Dynamically populated with agent names -1-> -->
<!-- </select> -->
<!-- </div> -->
<!-- <div class="mb-3"> -->
<!-- <label for="command" class="form-label">Command</label> -->
<!-- <input type="text" class="form-control" id="command" name="command" placeholder="Enter command" required> -->
<!-- </div> -->
<!-- <button type="submit" class="btn btn-primary">Execute</button> -->
<!-- </form> -->
<!-- <div id="commandOutput" class="mt-3 p-2 border">Here goes the output...</div> -->
<!-- </div> -->
<!-- <!-1- Add Agent Form -1-> -->
<!-- <button class="btn btn-primary mt-3" data-bs-toggle="collapse" data-bs-target="#addAgentForm">Add Agent</button> -->
<!-- <div id="addAgentForm" class="collapse mt-2"> -->
<!-- <form hx-post="/agents" hx-target="#agentList" hx-swap="innerHTML"> -->
<!-- <div class="mb-3"> -->
<!-- <label for="agentId" class="form-label">Agent Id</label> -->
<!-- <input type="text" class="form-control" id="agentId" name="agentId" required> -->
<!-- </div> -->
<!-- <div class="mb-3"> -->
<!-- <label for="agentName" class="form-label">Agent Name</label> -->
<!-- <input type="text" class="form-control" id="agentNameInput" name="agentName" required> -->
<!-- </div> -->
<!-- <div class="mb-3"> -->
<!-- <label for="IPv4Address" class="form-label">IPv4 Address</label> -->
<!-- <input type="text" class="form-control" id="IPv4Address" name="IPv4Address" required> -->
<!-- </div> -->
<!-- <div class="mb-3"> -->
<!-- <label for="initialContact" class="form-label">Initial Contact</label> -->
<!-- <input type="datetime-local" class="form-control" id="initialContact" name="initialContact" required> -->
<!-- </div> -->
<!-- <div class="mb-3"> -->
<!-- <label for="lastContact" class="form-label">Last Contact</label> -->
<!-- <input type="datetime-local" class="form-control" id="lastContact" name="lastContact" required> -->
<!-- </div> -->
<!-- <button type="submit" class="btn btn-success">Add Agent</button> -->
<!-- </form> -->
<!-- </div> -->
<!-- </div> -->
<!-- <!-1- Agent Details -1-> -->
<!-- <div class="col" id="agentDetails"> -->
<!-- <h3>Details</h3> -->
<!-- <p>Select an agent to view details.</p> -->
<!-- </div> -->
<!-- </div> -->
<!-- </div> -->
<!-- </body> -->
<!-- </html> -->
<!DOCTYPE html>
<html lang="en">
<head>
@ -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);
}
};

View File

@ -2,6 +2,7 @@
<h2>Agent Details</h2>
<p>ID: {{.AgentID}}</p>
<p>Name: {{.AgentName}}</p>
<p>Type: {{.AgentType}}</p>
<p>Initial Contact: {{.InitialContact}}</p>
<p>Last Contact: {{.LastContact}}</p>
<button hx-get="/agents" hx-target="#agentList" hx-swap="innerHTML">Back to List</button>

View File

@ -3,6 +3,7 @@
<tr>
<th>ID</th>
<th>Name</th>
<th>Type</th>
<th>IPv4 Address</th>
<th>Initial Contact</th>
<th>Last Contact</th>
@ -14,6 +15,7 @@
<tr>
<td>{{.AgentID}}</td>
<td>{{.AgentName}}</td>
<td>{{.AgentType}}</td>
<td>{{.IPv4Address}}</td>
<td>{{.InitialContact}}</td>
<td>{{.LastContact}}</td>