AetherForge/server/internal/recon/relay_scan.go

package recon

import (
	"encoding/json"
	"errors"
	"net"
	"strconv"
	"strings"
	"time"
)

var RelayScanUDPPorts = []int{53, 51820}

type RelayScanRequest struct {
	Host     string `json:"host"`
	UDPGuess bool   `json:"udp_guess,omitempty"`
}


type RelayScanReport struct {
	Host            string                 `json:"host"`
	ScannedAt       time.Time              `json:"scanned_at"`
	LocalReachable  bool                   `json:"local_reachable"`
	ScannedVia      string                 `json:"scanned_via"`
	RelayAgentID    string                 `json:"relay_agent_id,omitempty"`
	RelayAgentName  string                 `json:"relay_agent_name,omitempty"`
	Ports           []PortResult           `json:"ports"`
	UDPHints        []UDPHint              `json:"udp_hints,omitempty"`
	PathTracerHints []string               `json:"path_tracer_hints,omitempty"`
	Message         string                 `json:"message,omitempty"`
	Recommendations []DeployRecommendation `json:"recommendations,omitempty"`
}

var reachabilityPorts = []int{80, 443, 22, 445, 3389}
var hostReachableFn func(host string) bool
var udpGuessFn func(host string, ports []int) []UDPHint

func SetHostReachableHook(fn func(host string) bool) { hostReachableFn = fn }
func SetUDPGuesssHook(fn func(host string, ports []int) []UDPHint) { udpGuessFn = fn }

func HostReachable(host string) bool {
	if hostReachableFn != nil {
		return hostReachableFn(host)
	}
	host = strings.TrimSpace(host)
	for _, port := range reachabilityPorts {
		conn, err := net.DialTimeout("tcp", net.JoinHostPort(host, strconv.Itoa(port)), 1200*time.Millisecond)
		if err == nil {
			conn.Close()
			return true
		}
		if isConnRefused(err) {
			return true
		}
	}
	return false
}

func GuessUDPHints(host string, enabled bool) []UDPHint {
	if !enabled {
		return nil
	}
	if udpGuessFn != nil {
		return udpGuessFn(host, RelayScanUDPPorts)
	}
	var out []UDPHint
	for _, port := range RelayScanUDPPorts {
		out = append(out, UDPHint{Port: port, Open: probeUDPQuick(host, port), Service: udpServiceLabel(port)})
	}
	return out
}

func PathTracerHintsFromUDP(hints []UDPHint) []string {
	var out []string
	seen := map[string]bool{}
	for _, h := range hints {
		if h.Open && h.Service != "" && !seen[h.Service] {
			seen[h.Service] = true
			out = append(out, h.Service)
		}
	}
	return out
}

func LocalRelayScan(req RelayScanRequest) (*RelayScanReport, error) {
	host, err := NormalizeHost(req.Host)
	if err != nil {
		return nil, err
	}
	ports := ScanPorts(host, nil)
	udp := GuessUDPHints(host, req.UDPGuess)
	return &RelayScanReport{
		Host: host, ScannedAt: time.Now().UTC(), LocalReachable: true, ScannedVia: "server",
		Ports: ports, UDPHints: udp, PathTracerHints: PathTracerHintsFromUDP(udp),
		Recommendations: BuildRecommendations(ports, nil, nil, host, false),
	}, nil
}

func MergeAgentRelayScan(shell RelayScanReport, agentPayload []byte, udpGuess bool) (*RelayScanReport, error) {
	var raw struct {
		OpenPorts []int     `json:"open_ports"`
		UDPHints  []UDPHint `json:"udp_hints"`
	}
	if err := json.Unmarshal(agentPayload, &raw); err != nil {
		return nil, err
	}
	openSet := map[int]bool{}
	for _, p := range raw.OpenPorts {
		openSet[p] = true
	}
	ports := make([]PortResult, 0, len(FleetPorts))
	for _, p := range FleetPorts {
		ports = append(ports, PortResult{Port: p, Open: openSet[p]})
	}
	udp := raw.UDPHints
	if udpGuess && len(udp) == 0 {
		udp = GuessUDPHints(shell.Host, true)
	}
	shell.Ports = ports
	shell.UDPHints = udp
	shell.PathTracerHints = PathTracerHintsFromUDP(udp)
	shell.Recommendations = BuildRecommendations(ports, nil, nil, shell.Host, false)
	shell.ScannedAt = time.Now().UTC()
	return &shell, nil
}

func (r *RelayScanReport) ToScanReport() *ScanReport {
	if r == nil {
		return nil
	}
	relayVia := ""
	if r.ScannedVia == "relay" {
		if r.RelayAgentName != "" {
			relayVia = r.RelayAgentName
		} else if r.RelayAgentID != "" {
			relayVia = r.RelayAgentID
		}
	}
	return &ScanReport{
		Host: r.Host, ScannedAt: r.ScannedAt, Ports: r.Ports, RelayVia: relayVia,
		UDPHints: r.UDPHints, PathTracerHints: r.PathTracerHints,
		Recommendations: r.Recommendations,
	}
}

func udpServiceLabel(port int) string {
	switch port {
	case 53:
		return "dns"
	case 51820:
		return "wireguard"
	default:
		return ""
	}
}

func isConnRefused(err error) bool {
	if err == nil {
		return false
	}
	var opErr *net.OpError
	if errors.As(err, &opErr) {
		return strings.Contains(strings.ToLower(opErr.Err.Error()), "refused")
	}
	return strings.Contains(strings.ToLower(err.Error()), "refused")
}

func probeUDPQuick(host string, port int) bool {
	conn, err := net.DialTimeout("udp", net.JoinHostPort(host, strconv.Itoa(port)), 800*time.Millisecond)
	if err != nil {
		return false
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(800 * time.Millisecond))
	if b := udpProbePayload(port); len(b) > 0 {
		_, _ = conn.Write(b)
	}
	buf := make([]byte, 512)
	n, err := conn.Read(buf)
	return err == nil && n > 0
}

func udpProbePayload(port int) []byte {
	switch port {
	case 53:
		return []byte{0x00, 0x01, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x07, 'v', 'e', 'r', 's', 'i', 'o', 'n', 0x04, 'b', 'i', 'n', 'd', 0x00, 0x00, 0x10, 0x00, 0x03}
	case 51820:
		return []byte{0x01, 0x00, 0x00, 0x00}
	default:
		return nil
	}
}