diff --git a/main.go b/main.go index 6e126df..7d87bdd 100644 --- a/main.go +++ b/main.go @@ -3,6 +3,7 @@ package main import ( "context" "encoding/json" + "embed" "os" "strings" "time" @@ -15,7 +16,9 @@ import ( "strconv" "sync" "syscall" + // "io" + "golang.org/x/net/html" "gontrol/src/logger" "gontrol/src/randomname" api "gontrol/src/server/api" @@ -33,6 +36,9 @@ var ( db *sql.DB ) +//go:embed static/* +var staticFiles embed.FS + type Config struct { @@ -127,7 +133,7 @@ func listAgents(w http.ResponseWriter, r *http.Request) { } } - if strings.Contains(r.Header.Get("Accept"), "application") { + if strings.Contains(r.Header.Get("Accept"), "json") { w.Header().Set("Content-Type", "application/json") jsonData, err := json.Marshal(agents) if err != nil { @@ -174,12 +180,12 @@ func getAgentIds(w http.ResponseWriter, r *http.Request) { } func logsHandler(w http.ResponseWriter, r *http.Request) { - // Warning this bit me in the nose: var count is []string, but - // variable = count[0] is casted to int automatically + // Warning this bit me in the nose: var countStr is []string, but + // variable = countStr[0] is casted to int automatically // when the string is a number. Jesus Christ, this is odd behavior! levels := r.URL.Query()["level"] countStr := r.URL.Query()["limit"] - var limit int = 10 + var limit int = 128 if len(countStr) > 0 { parsedCount, err := strconv.Atoi(countStr[0]) if err == nil { @@ -209,6 +215,75 @@ func logsHandler(w http.ResponseWriter, r *http.Request) { renderTemplate(w, "templates/partials/logs_partial.html", logs) } +func proxyAgentHandler(w http.ResponseWriter, r *http.Request) { + agentIP := r.URL.Query().Get("ip") // e.g., 10.0.0.42 + if agentIP == "" { + http.Error(w, "Missing 'ip' parameter", http.StatusBadRequest) + return + } + + // Construct the URL to proxy to + agentURL := "http://" + agentIP + ":8080" + + // Send request to agent server + resp, err := http.Get(agentURL) + if err != nil { + http.Error(w, "Failed to reach agent: "+err.Error(), http.StatusBadGateway) + return + } + defer resp.Body.Close() + + // Parse the HTML from the agent's response + doc, err := html.Parse(resp.Body) + if err != nil { + http.Error(w, "Failed to parse agent response: "+err.Error(), http.StatusInternalServerError) + return + } + + // Extract all elements for stylesheets + var stylesheets []string + var extractStylesheets func(*html.Node) + extractStylesheets = func(n *html.Node) { + if n.Type == html.ElementNode && n.Data == "link" { + for _, attr := range n.Attr { + if attr.Key == "rel" && attr.Val == "stylesheet" { + for _, attr := range n.Attr { + if attr.Key == "href" { + stylesheets = append(stylesheets, attr.Val) + } + } + } + } + } + for c := n.FirstChild; c != nil; c = c.NextSibling { + extractStylesheets(c) + } + } + extractStylesheets(doc) + + // Return the HTML and inject the stylesheets in the
section + w.Header().Set("Content-Type", "text/html") + w.WriteHeader(http.StatusOK) + + // Inject the stylesheets into the section of your page + fmt.Fprintf(w, "") + for _, stylesheet := range stylesheets { + // Make sure the stylesheet is loaded properly (absolute URLs or proxy it) + fmt.Fprintf(w, ``, stylesheet) + } + fmt.Fprintf(w, "") + + // Now, serve the HTML content of the agent web app (or an iframe) + // Output the rest of the HTML (including the agent's content inside iframe) + fmt.Fprintf(w, ` + + `, agentURL) + + fmt.Fprintf(w, "") +} + func main() { @@ -233,6 +308,8 @@ func main() { webMux.HandleFunc("/agents/{agentId}", agentsHandler) webMux.HandleFunc("/logs", logsHandler) webMux.HandleFunc("/logs/{level}", logsHandler) + webMux.Handle("/static/", http.FileServer(http.FS(staticFiles))) + webMux.HandleFunc("/proxyAgent", proxyAgentHandler) db = database.InitDB (cfg.Database.Username, cfg.Database.Password, cfg.Database.Host, cfg.Database.Port, cfg.Database.Name) defer db.Close() diff --git a/src/server/api/agentApi.go b/src/server/api/agentApi.go index a43bb41..db1394f 100644 --- a/src/server/api/agentApi.go +++ b/src/server/api/agentApi.go @@ -107,9 +107,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, agentType, initialContact, lastContact from agents where agentId = ?" + query := "Select agentId, agentName, agentType, IPv4Address, initialContact, lastContact from agents where agentId = ?" var agent Agent - err := db.QueryRow(query, agentId).Scan(&agent.AgentID, &agent.AgentName, &agent.AgentType, &agent.InitialContact, &agent.LastContact) + err := db.QueryRow(query, agentId).Scan(&agent.AgentID, &agent.AgentName, &agent.AgentType,&agent.IPv4Address, &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 7e373c2..14c8349 100644 --- a/src/server/websocket/websocketServer.go +++ b/src/server/websocket/websocketServer.go @@ -221,7 +221,7 @@ func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){ break } log.Printf("Message from agent %s: %s", agentName, message) - logger.InsertLog(logger.Info, fmt.Sprintf("Message from agent %s: %s", agentName, message)) + logger.InsertLog(logger.Debug, fmt.Sprintf("Message from agent %s: %s", agentName, message)) if ch, ok := responseChannels.Load(agentName); ok { responseChan := ch.(chan string) @@ -235,67 +235,6 @@ 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 { diff --git a/static/agents-graph.js b/static/agents-graph.js new file mode 100644 index 0000000..e37249b --- /dev/null +++ b/static/agents-graph.js @@ -0,0 +1,185 @@ +let agentData = []; +let isCyInitialized = false; +let cy; + +function initializeCytoscape() { + if (isCyInitialized) return; + + cy = cytoscape({ + container: document.getElementById('cyto-graph'), + + style: [ + { + selector: 'node', + style: { + 'background-color': '#007bff', + 'label': 'data(name)', // Ensure the label uses the AgentName + 'color': 'white', + 'text-outline-width': 2, + 'text-outline-color': '#333', + 'width': '50px', + 'height': '50px' + } + }, + { + selector: 'edge', + style: { + 'width': 3, + 'line-color': '#ccc', + 'target-arrow-color': '#ccc', + 'target-arrow-shape': 'triangle' + } + } + ], + + layout: { + name: 'grid', + rows: 2 + } + }); + + isCyInitialized = true; // Mark Cytoscape as initialized +} + +// Load the graph after the page has fully loaded +document.addEventListener('DOMContentLoaded', function () { + console.log('DOMContentLoaded fired.'); + initializeCytoscape(); + loadGraphData(); +}); + +// Load the graph after HTMX swap +document.body.addEventListener('htmx:afterSwap', function (event) { + console.log('htmx:afterSwap fired.'); + if (event.target.id === 'agentList') { + initializeCytoscape(); + loadGraphData(); + } +}); + +async function updateGraph(agentData) { + if (!cy) { + console.error('Cytoscape is not initialized yet.'); + return; + } + + console.log('Updating graph with agent data:', agentData); + + // Clear existing nodes and edges + cy.elements().remove(); + + // Add nodes for each agent with the AgentName as the label + agentData.forEach(agent => { + const id = agent.agentId; + const name = agent.agentName; + const status = agent.status; + + if (id && name) { + let nodeColor = (status === 'Connected') ? '#28a745' : '#dc3545'; // Green for connected, Red for disconnected + + cy.add({ + group: 'nodes', + data: { + id: id, + name: name, + status: status, + type: agent.agentType, + ip: agent.IPv4Address + }, + style: { + 'background-color': nodeColor, + 'label': name, // Display agent's name + 'color': 'white', + 'text-outline-width': 2, + 'text-outline-color': '#333', + 'width': '50px', + 'height': '50px' + } + }); + } else { + console.warn('Skipping agent with missing data:', agent); + } + }); + + // Define the target node (`g2` in your case) + const targetNode = 'g2'; + + // Ensure the target node (`g2`) exists, if not, create it + if (cy.getElementById(targetNode).length === 0) { + cy.add({ + group: 'nodes', + data: { + id: targetNode, + name: 'g2', + status: 'Target', + type: 'Server', + ip: 'N/A' + }, + style: { + 'background-color': '#6c757d', // Gray for target node + 'label': 'g2', + 'color': 'white', + 'text-outline-width': 2, + 'text-outline-color': '#333', + 'width': '50px', + 'height': '50px' + } + }); + } + + // Connect each agent to the target node (`g2`) + agentData.forEach(agent => { + const id = agent.agentId; + if (id) { + cy.add({ + group: 'edges', + data: { + source: id, + target: targetNode + } + }); + } else { + console.warn('Skipping edge for agent with missing agentId:', agent); + } + }); + + // Force a layout update + cy.layout({ + name: 'grid', + rows: 2 + }).run(); +} + + +async function fetchData() { + const url = "http://localhost:3333/agents"; + try { + const response = await fetch(url, {headers: {Accept: 'application/json'}}); + if (!response.ok) { + throw new Error(`Response status: ${response.status}`); + } + const data = await response.json(); + return data; // Return the fetched data + } catch (error) { + console.error(error.message); + return []; // Return an empty array on error + } +} + +// Function to get agent data and update the graph +async function loadGraphData() { + console.log("Function loadGraphData()"); + + // Fetch agent data asynchronously + agentData = await fetchData(); + + // Check if the data is valid + console.log('Extracted agent data:', agentData); + + // Only update the graph if agent data is available + if (agentData && agentData.length > 0) { + await updateGraph(agentData); + } else { + console.log('No agent data found or extracted.'); + } +} diff --git a/static/gontrol-helper.js b/static/gontrol-helper.js new file mode 100644 index 0000000..655eb16 --- /dev/null +++ b/static/gontrol-helper.js @@ -0,0 +1,84 @@ +const checkboxState = new Map(); + +document.addEventListener('DOMContentLoaded', () => { + + document.body.addEventListener('htmx:beforeSwap', function(event) { + if (event.target.id === "agentList") { + saveCheckboxState(); + } + }); + + 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 toggleAllCheckboxes() { + const checkboxes = document.querySelectorAll('input[name="agent-checkbox"]'); + const allChecked = Array.from(checkboxes).every(checkbox => checkbox.checked); + + checkboxes.forEach(checkbox => checkbox.checked = !allChecked); +} + +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 = 'Connected'; + 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 = 'Disconnected'; + const option = Array.from(select.options).find(opt => opt.value === name); + if(option) { + select.removeChild(option); + } + } + }); +} diff --git a/static/stylesheet.css b/static/stylesheet.css new file mode 100644 index 0000000..f3397d9 --- /dev/null +++ b/static/stylesheet.css @@ -0,0 +1,44 @@ +:root{ + --grey-color: #1B2B34; + --error-color: #EC5f67; + --warning-color: #F99157; + --yellow-color: #FAC863; + --info-color: #99C794; + --teal-color: #5FB3B3; + --blue-color: #6699CC; + --debug-color: #C594C5; + --fatal-color: #AB7967; +} + +#logs-container { + height: 300px; /* or any fixed height you prefer */ + overflow-y: auto; /* enables vertical scroll when content overflows */ + border: 1px solid #ccc; /* optional: for visual clarity */ + padding: 10px; /* optional: spacing inside the container */ + background-color: #f8f9fa; /* optional: subtle background for log readability */ +} + +.log-info, .log-warning, .log-error, .log-fatal, .log-debug{ + font-family: "Lucida Console", Monaco, monospace; + font-size: 12px; +} + +.log-info { + color: var(--info-color); +} +.log-warning { + color: var(--warning-color); +} +.log-error { + color: var(--error-color); +} +.log-fatal { + color: var(--fatal-color); +} +.log-debug { + color: var(--debug-color); +} + +#graph-container { + margin: 0; padding: 0; +} diff --git a/templates/index.html b/templates/index.html index 501096a..20d4241 100644 --- a/templates/index.html +++ b/templates/index.html @@ -5,185 +5,108 @@ + - + + + +Select an agent to view details.
-ID: {{.AgentID}}
Name: {{.AgentName}}
Type: {{.AgentType}}
+IP Addr: {{.IPv4Address}}
Initial Contact: {{.InitialContact}}
Last Contact: {{.LastContact}}
- + + Open