AetherForge/server/internal/api/recon_relay_scan_test.go

package api

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

	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/recon"
)

func TestReconRelayScanLocalPath(t *testing.T) {
	recon.SetHostReachableHook(func(host string) bool { return true })
	t.Cleanup(func() { recon.SetHostReachableHook(nil) })
	recon.SetPortDialHook(func(host string, port int, _ time.Duration) bool { return port == 22 })
	t.Cleanup(func() { recon.SetPortDialHook(nil) })

	h := NewReconHandler(nil, nil)
	body, _ := json.Marshal(map[string]interface{}{"host": "10.0.0.10"})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/recon/relay-scan", bytes.NewReader(body))
	w := httptest.NewRecorder()
	h.RelayScan(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
	}
	var resp map[string]interface{}
	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
		t.Fatal(err)
	}
	if resp["ok"] != true {
		t.Fatalf("resp=%v", resp)
	}
}

func TestReconRelayScanNoRelayFallback(t *testing.T) {
	recon.SetHostReachableHook(func(host string) bool { return false })
	t.Cleanup(func() { recon.SetHostReachableHook(nil) })

	h := NewReconHandler(nil, NewWSHub(nil))
	body, _ := json.Marshal(map[string]interface{}{"host": "10.99.1.50"})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/recon/relay-scan", bytes.NewReader(body))
	w := httptest.NewRecorder()
	h.RelayScan(w, req)
	var resp map[string]interface{}
	_ = json.Unmarshal(w.Body.Bytes(), &resp)
	if resp["ok"] != false || resp["error"] == nil {
		t.Fatalf("resp=%v", resp)
	}
}

func TestReconRelayScanFleetRelay(t *testing.T) {
	recon.SetHostReachableHook(func(host string) bool { return false })
	t.Cleanup(func() { recon.SetHostReachableHook(nil) })

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

	hub := NewWSHub(database)
	agentID := "relay-seed"
	conn := connectTestAgentWithIP(t, hub, agentID, "10.42.1.50")

	done := make(chan struct{})
	go func() {
		defer close(done)
		deadline := time.Now().Add(5 * time.Second)
		for time.Now().Before(deadline) {
			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 {
				return
			}
			if payload["action"] != "recon_relay_scan" {
				continue
			}
			result := `{"open_ports":[22,445],"udp_hints":[{"port":53,"open":true,"service":"dns"}]}`
			cmdPayload, _ := json.Marshal(map[string]interface{}{
				"action": "recon_relay_scan", "success": true, "message": result,
			})
			_ = conn.WriteJSON(Message{Type: "command_result", Payload: cmdPayload})
			return
		}
	}()

	h := NewReconHandler(database, hub)
	body, _ := json.Marshal(map[string]interface{}{"host": "10.42.1.100", "udp_guess": true})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/recon/relay-scan", bytes.NewReader(body))
	w := httptest.NewRecorder()
	h.RelayScan(w, req)

	select {
	case <-done:
	case <-time.After(6 * time.Second):
		t.Fatal("agent did not receive relay command")
	}

	var resp struct {
		OK      bool                   `json:"ok"`
		AgentID string                 `json:"agent_id"`
		Report  map[string]interface{} `json:"report"`
	}
	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
		t.Fatal(err)
	}
	if !resp.OK || resp.AgentID != agentID {
		t.Fatalf("resp=%+v", resp)
	}
	hints, _ := resp.Report["path_tracer_hints"].([]interface{})
	if len(hints) != 1 || hints[0] != "dns" {
		t.Fatalf("report=%v", resp.Report)
	}
}