AetherForge/server/internal/api/pathtracer_discover_test.go

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