added multi agent execution
This commit is contained in:
Stefan Etringer
parent
2056479224
commit
1ce6d2e676
18
main.go
18
main.go
|
|
@ -6,7 +6,7 @@ import (
|
|||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"slices"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"html/template"
|
||||
|
|
@ -118,17 +118,11 @@ func listAgents(w http.ResponseWriter, r *http.Request) {
|
|||
agents, err := api.GetAgents(db)
|
||||
currentAgents := getAgentsStatus()
|
||||
|
||||
for _, currAgent := range currentAgents {
|
||||
for i, agent := range agents {
|
||||
if currAgent == agent.AgentName {
|
||||
// log.Printf("%s online", agent.AgentName)
|
||||
// logger.InsertLog(logger.Debug, fmt.Sprintf("%s online after page refresh", agent.AgentName))
|
||||
// agents[i].Status = fmt.Sprint("<span class=\"badge bg-success\">Connected</span>")
|
||||
agents[i].Status = "Connected"
|
||||
} else {
|
||||
// agent.Status = fmt.Sprintf("<span class=\"badge bg-danger\">Disconnected</span>")
|
||||
agents[i].Status = "Disconnected"
|
||||
}
|
||||
for i := range agents {
|
||||
if slices.Contains(currentAgents, agents[i].AgentName) {
|
||||
agents[i].Status = "Connected"
|
||||
} else {
|
||||
agents[i].Status = "Disconnected"
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -235,6 +235,67 @@ type Message struct {
|
|||
Payload string `json:"payload"`
|
||||
}
|
||||
|
||||
// var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Request){
|
||||
// err := r.ParseForm()
|
||||
// if err != nil {
|
||||
// http.Error(w, "Invalid form data", http.StatusBadRequest)
|
||||
// logger.InsertLog(logger.Info, "Invalid form data")
|
||||
// return
|
||||
// }
|
||||
|
||||
// agentName := r.FormValue("agentName")
|
||||
// command := r.FormValue("command")
|
||||
|
||||
// agentSocketsMutex.Lock()
|
||||
// conn, ok := agentSockets[agentName]
|
||||
// agentSocketsMutex.Unlock()
|
||||
|
||||
// if !ok {
|
||||
// http.Error(w, "Agent not connected", http.StatusNotFound)
|
||||
// logger.InsertLog(logger.Info, "Agent not connected")
|
||||
// return
|
||||
// }
|
||||
|
||||
// responseChan := make(chan string, 1)
|
||||
// responseChannels.Store(agentName, responseChan)
|
||||
// defer responseChannels.Delete(agentName)
|
||||
|
||||
// message := Message {
|
||||
// Type: "command",
|
||||
// Payload: command,
|
||||
// }
|
||||
|
||||
// messageBytes, _ := json.Marshal(message)
|
||||
|
||||
// err = conn.WriteMessage(websocket.TextMessage, messageBytes)
|
||||
// if err != nil {
|
||||
// http.Error(w, "Failed to send command to the agent", http.StatusInternalServerError)
|
||||
// logger.InsertLog(logger.Error, "Failed to send command to the agent")
|
||||
// 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)
|
||||
// logger.InsertLog(logger.Error, "Invalid response structure")
|
||||
// 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 response timed out", http.StatusGatewayTimeout)
|
||||
// logger.InsertLog(logger.Info, "Agent response timed out")
|
||||
// }
|
||||
// }
|
||||
|
||||
var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Request){
|
||||
err := r.ParseForm()
|
||||
if err != nil {
|
||||
|
|
@ -243,59 +304,103 @@ var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Reques
|
|||
return
|
||||
}
|
||||
|
||||
agentName := r.FormValue("agentName")
|
||||
agentNameStr := r.FormValue("agentNames")
|
||||
var agentNames []string
|
||||
|
||||
if agentNameStr != "" {
|
||||
agentNames = strings.Split(agentNameStr, ",")
|
||||
} else {
|
||||
agentName := r.FormValue("agentName")
|
||||
if agentName != "" {
|
||||
agentNames = []string{agentName}
|
||||
}
|
||||
}
|
||||
|
||||
command := r.FormValue("command")
|
||||
|
||||
agentSocketsMutex.Lock()
|
||||
conn, ok := agentSockets[agentName]
|
||||
agentSocketsMutex.Unlock()
|
||||
|
||||
if !ok {
|
||||
http.Error(w, "Agent not connected", http.StatusNotFound)
|
||||
logger.InsertLog(logger.Info, "Agent not connected")
|
||||
if len(agentNames) == 0 || command == "" {
|
||||
http.Error(w, "Missing agent or command", http.StatusBadRequest)
|
||||
logger.InsertLog(logger.Error, "Missing agent or command")
|
||||
return
|
||||
}
|
||||
|
||||
responseChan := make(chan string, 1)
|
||||
responseChannels.Store(agentName, responseChan)
|
||||
defer responseChannels.Delete(agentName)
|
||||
|
||||
message := Message {
|
||||
Type: "command",
|
||||
Payload: command,
|
||||
type result struct {
|
||||
AgentName string
|
||||
Type string
|
||||
Payload string
|
||||
Err error
|
||||
}
|
||||
|
||||
messageBytes, _ := json.Marshal(message)
|
||||
resultsChan := make(chan result, len(agentNames))
|
||||
|
||||
err = conn.WriteMessage(websocket.TextMessage, messageBytes)
|
||||
if err != nil {
|
||||
http.Error(w, "Failed to send command to the agent", http.StatusInternalServerError)
|
||||
logger.InsertLog(logger.Error, "Failed to send command to the agent")
|
||||
return
|
||||
for _, agentName := range agentNames {
|
||||
agentName := strings.TrimSpace(agentName)
|
||||
|
||||
go func(agent string) {
|
||||
agentSocketsMutex.Lock()
|
||||
conn, ok := agentSockets[agentName]
|
||||
agentSocketsMutex.Unlock()
|
||||
|
||||
if !ok {
|
||||
resultsChan <- result{AgentName: agent, Err: fmt.Errorf("Agent not connected")}
|
||||
return
|
||||
}
|
||||
|
||||
responseChan := make(chan string, 1)
|
||||
responseChannels.Store(agent, responseChan)
|
||||
defer responseChannels.Delete(agent)
|
||||
|
||||
msg := Message {
|
||||
Type: "command",
|
||||
Payload: command,
|
||||
}
|
||||
|
||||
msgBytes, _ := json.Marshal(msg)
|
||||
err := conn.WriteMessage(websocket.TextMessage, msgBytes)
|
||||
if err != nil {
|
||||
resultsChan <- result{AgentName: agent, Err: fmt.Errorf("Send failed")}
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case resp := <- responseChan:
|
||||
var parsed map[string]string
|
||||
if err:= json.Unmarshal([]byte(resp), &parsed); err != nil {
|
||||
resultsChan <- result{AgentName: agent, Err: fmt.Errorf("Invalid response")}
|
||||
return
|
||||
}
|
||||
|
||||
payload, ok := parsed["payload"]
|
||||
if !ok {
|
||||
resultsChan <- result{AgentName: agent, Err: fmt.Errorf("No payload")}
|
||||
return
|
||||
}
|
||||
|
||||
resultsChan <- result{AgentName: agent, Payload: payload}
|
||||
case <-time.After(10 * time.Second):
|
||||
resultsChan <- result{AgentName: agent, Err: fmt.Errorf("Timeout")}
|
||||
}
|
||||
} (agentName)
|
||||
}
|
||||
|
||||
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)
|
||||
logger.InsertLog(logger.Error, "Invalid response structure")
|
||||
return
|
||||
var combined strings.Builder
|
||||
for i := 0; i < len(agentNames); i++ {
|
||||
res := <- resultsChan
|
||||
if res.Err != nil {
|
||||
combined.WriteString(fmt.Sprintf("[%s] ERROR: %s\n", res.AgentName, res.Err.Error()))
|
||||
} else {
|
||||
combined.WriteString(fmt.Sprintf("[%s] %s\n", res.AgentName, res.Payload))
|
||||
}
|
||||
}
|
||||
|
||||
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 response timed out", http.StatusGatewayTimeout)
|
||||
logger.InsertLog(logger.Info, "Agent response timed out")
|
||||
}
|
||||
w.Write([]byte(combined.String()))
|
||||
|
||||
|
||||
}
|
||||
|
||||
|
||||
func Server() (*http.Server) {
|
||||
webSocketHandler := webSocketHandler {
|
||||
upgrader: websocket.Upgrader{
|
||||
|
|
|
|||
|
|
@ -10,101 +10,84 @@
|
|||
<!-- <script src="https://code.jquery.com/jquery-3.7.1.min.js"></script> -->
|
||||
<title>g2: gommand & gontrol</title>
|
||||
<script>
|
||||
// Query Agents for the Dropdown Menu
|
||||
<!-- document.addEventListener('DOMContentLoaded', () => { -->
|
||||
<!-- 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)); -->
|
||||
<!-- }); -->
|
||||
<!-- // Query agents currently connected to the websocket and put status into the table -->
|
||||
<!-- const updateAgentStatuses = () => { -->
|
||||
<!-- fetch('http://localhost:5555/agentNames') -->
|
||||
<!-- .then(response => response.json()) -->
|
||||
<!-- .then(agentNames => { -->
|
||||
<!-- console.log("Agent names fetched:", agentNames); -->
|
||||
<!-- const tableRows = document.querySelectorAll('#agentList table tbody tr'); -->
|
||||
<!-- tableRows.forEach(row => { -->
|
||||
<!-- const nameCell = row.querySelector('td:nth-child(2)'); -->
|
||||
<!-- const statusCell = row.querySelector('td:nth-child(5)'); -->
|
||||
<!-- if (nameCell && statusCell) { -->
|
||||
<!-- const agentName = nameCell.textContent.trim(); -->
|
||||
<!-- if (agentNames.includes(agentName)) { -->
|
||||
<!-- statusCell.innerHTML = '<span class="badge bg-success">Connected</span>'; -->
|
||||
<!-- } else { -->
|
||||
<!-- statusCell.innerHTML = '<span class="badge bg-danger">Disconnected</span>'; -->
|
||||
<!-- } -->
|
||||
<!-- } -->
|
||||
<!-- }); -->
|
||||
<!-- }) -->
|
||||
<!-- .catch(error => console.error('Error fetching agent names:', error)); -->
|
||||
<!-- }; -->
|
||||
<!-- updateAgentStatuses(); -->
|
||||
<!-- setInterval(updateAgentStatuses, 5000); -->
|
||||
|
||||
<!-- document.body.addEventListener('htmx:afterSwap', function(evt) { -->
|
||||
<!-- if (evt.detail.xhr.status === 200) { -->
|
||||
<!-- const tableRows = document.querySelectorAll('#agentList table tbody tr'); -->
|
||||
<!-- tableRows.forEach(row => { -->
|
||||
<!-- const nameCell = row.querySelector('td:nth-child(2)'); -->
|
||||
<!-- const statusCell = row.querySelector('td:nth-child(5)'); -->
|
||||
<!-- if (nameCell && statusCell) { -->
|
||||
<!-- const agentName = nameCell.textContent.trim(); -->
|
||||
<!-- if ("Connected" === statusCell.innerHTML) { -->
|
||||
<!-- statusCell.innerHTML = '<span class="badge bg-success">Connected</span>'; -->
|
||||
<!-- } else { -->
|
||||
<!-- statusCell.innerHTML = '<span class="badge bg-danger">Disconnected</span>'; -->
|
||||
<!-- } -->
|
||||
<!-- } -->
|
||||
<!-- } -->
|
||||
<!-- } -->
|
||||
<!-- }); -->
|
||||
const checkboxState = new Map();
|
||||
|
||||
document.addEventListener('DOMContentLoaded', () => {
|
||||
document.body.addEventListener('htmx:afterSwap', function(event) {
|
||||
|
||||
document.body.addEventListener('htmx:beforeSwap', function(event) {
|
||||
if (event.target.id === "agentList") {
|
||||
updateAgentDropdown();
|
||||
saveCheckboxState();
|
||||
}
|
||||
});
|
||||
|
||||
function updateAgentDropdown() {
|
||||
const select = document.getElementById("agentName");
|
||||
const optionValues = Array.from(select.options).map(opt => opt.value);
|
||||
const rows = document.querySelectorAll("#agentList tbody tr");
|
||||
|
||||
rows.forEach(row => {
|
||||
const status = row.cells[4].textContent.trim();
|
||||
const name = row.cells[1].textContent.trim();
|
||||
|
||||
if (status === "Connected") {
|
||||
row.cells[4].innerHTML = '<span class="badge bg-success">Connected</span>';
|
||||
const option = document.createElement("option");
|
||||
if (!(optionValues.includes(name))) {
|
||||
option.value = name;
|
||||
option.textContent = name;
|
||||
select.appendChild(option);
|
||||
}
|
||||
}
|
||||
|
||||
if (status === "Disconnected") {
|
||||
row.cells[4].innerHTML = '<span class="badge bg-danger">Disconnected</span>';
|
||||
const option = Array.from(select.options).find(opt => opt.value === name);
|
||||
if(option) {
|
||||
select.removeChild(option);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
document.body.addEventListener('htmx:afterSwap', function(event) {
|
||||
if (event.target.id === "agentList") {
|
||||
restoreCheckboxState();
|
||||
updateAgentDropdown();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
function prepareAgentNames(event) {
|
||||
const selected = Array.from(document.querySelectorAll('.agent-checkbox'))
|
||||
.filter(cb => cb.checked)
|
||||
.map(cb => cb.dataset.agentName);
|
||||
|
||||
const hiddenInput = document.getElementById('agentNamesInput');
|
||||
|
||||
if (selected.length > 0) {
|
||||
document.getElementById('agentName').removeAttribute('name');
|
||||
hiddenInput.value = selected.join(',');
|
||||
} else {
|
||||
document.getElementById('agentName').setAttribute('name', 'agentName');
|
||||
hiddenInput.value = '';
|
||||
}
|
||||
}
|
||||
|
||||
function saveCheckboxState() {
|
||||
document.querySelectorAll('.agent-checkbox').forEach((checkbox) => {
|
||||
checkboxState.set(checkbox.dataset.agentName, checkbox.checked);
|
||||
});
|
||||
}
|
||||
|
||||
function restoreCheckboxState() {
|
||||
document.querySelectorAll('.agent-checkbox').forEach((checkbox) => {
|
||||
const state = checkboxState.get(checkbox.dataset.agentName);
|
||||
if (state !== undefined) {
|
||||
checkbox.checked = state;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
function updateAgentDropdown() {
|
||||
const select = document.getElementById("agentName");
|
||||
const optionValues = Array.from(select.options).map(opt => opt.value);
|
||||
const rows = document.querySelectorAll("#agentList tbody tr");
|
||||
|
||||
rows.forEach(row => {
|
||||
const status = row.cells[4].textContent.trim();
|
||||
const name = row.cells[1].textContent.trim();
|
||||
|
||||
if (status === "Connected") {
|
||||
row.cells[4].innerHTML = '<span class="badge bg-success">Connected</span>';
|
||||
const option = document.createElement("option");
|
||||
if (!(optionValues.includes(name))) {
|
||||
option.value = name;
|
||||
option.textContent = name;
|
||||
select.appendChild(option);
|
||||
}
|
||||
}
|
||||
|
||||
if (status === "Disconnected") {
|
||||
row.cells[4].innerHTML = '<span class="badge bg-danger">Disconnected</span>';
|
||||
const option = Array.from(select.options).find(opt => opt.value === name);
|
||||
if(option) {
|
||||
select.removeChild(option);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
</script>
|
||||
<style>
|
||||
.log-info {
|
||||
|
|
@ -137,11 +120,11 @@
|
|||
<!-- Agent Commands -->
|
||||
<div id="agentCommands">
|
||||
<h3>Command Execution</h3>
|
||||
<form hx-post="http://localhost:5555/executeCommand" hx-target="#commandOutput" hx-encoding="application/x-www-form-urlencoded" hx-swap="innerHTML">
|
||||
<form hx-post="http://localhost:5555/executeCommand" hx-target="#commandOutput" hx-encoding="application/x-www-form-urlencoded" hx-swap="innerHTML" onsubmit="prepareAgentNames(event)">
|
||||
<div class="mb-3">
|
||||
<label for="agentName" class="form-label">Agent Name</label>
|
||||
<!-- <select class="form-select" id="agentName" name="agentName" required> -->
|
||||
<select id="agentName" class="form-select" name="agentName" hx-on="htmx:afterSwap:updateAgentDropdown" required>
|
||||
<select id="agentName" class="form-select" name="agentName" hx-on="htmx:afterSwap:updateAgentDropdown">
|
||||
<option value="" disabled selected>Select an Agent</option>
|
||||
<!-- Dynamically populated with agent names -->
|
||||
</select>
|
||||
|
|
@ -151,39 +134,12 @@
|
|||
<input type="text" class="form-control" id="command" name="command" placeholder="Enter command" required>
|
||||
</div>
|
||||
<button type="submit" class="btn btn-primary">Execute</button>
|
||||
<!-- Hidden checkbox form !-->
|
||||
<input type="hidden" name="agentNames" id="agentNamesInput">
|
||||
</form>
|
||||
<pre id="commandOutput"></pre>
|
||||
</div>
|
||||
|
||||
<!-- Add Agent Form -->
|
||||
<!-- <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="agentName" 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> -->
|
||||
|
||||
<!-- Logs Section -->
|
||||
<h3>Logs</h3>
|
||||
<div id="logs-container" hx-get="/logs" hx-target="#logs-container" hx-swap="innerHTML" hx-trigger="every 3s">
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@
|
|||
<!-- <td><span class="badge bg-danger">Disconnected</span></td> -->
|
||||
<td>{{.Status}}</td>
|
||||
<td>
|
||||
<input type="checkbox" class="agent-checkbox" data-agent-name="{{.AgentName}}">
|
||||
<button class="btn btn-warning" hx-get="/agents/{{.AgentID}}" hx-target="#agentDetails" hx-swap="innerHTML">View</button>
|
||||
|
||||
<button class="btn btn-danger" hx-delete="/agents/{{.AgentID}}" hx-target="#agentList" hx-swap="innerHTML">Delete</button>
|
||||
|
|
|
|||
Loading…
Reference in New Issue