#!/usr/bin/env python3
"""Write recon batch 2 files, test, commit, and push in one shot."""
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]

FILES = {}

FILES["server/internal/recon/relay_scan.go"] = r'''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 UDPHint struct {
	Port    int    `json:"port"`
	Open    bool   `json:"open"`
	Service string `json:"service"`
}

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

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)
	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, Message: r.Message,
		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
	}
}
'''

FILES["server/internal/recon/relay_scan_test.go"] = r'''package recon

import "testing"

func TestHostReachableHook(t *testing.T) {
	SetHostReachableHook(func(host string) bool { return host == "10.0.0.5" })
	t.Cleanup(func() { SetHostReachableHook(nil) })
	if !HostReachable("10.0.0.5") || HostReachable("10.0.0.6") {
		t.Fatal("reachability hook mismatch")
	}
}

func TestMergeAgentRelayScanWireguardHint(t *testing.T) {
	payload := []byte(`{"open_ports":[22],"udp_hints":[{"port":51820,"open":true,"service":"wireguard"}]}`)
	merged, err := MergeAgentRelayScan(RelayScanReport{Host: "10.1.2.50", ScannedVia: "relay"}, payload, false)
	if err != nil || len(merged.PathTracerHints) != 1 || merged.PathTracerHints[0] != "wireguard" {
		t.Fatalf("merged=%+v err=%v", merged, err)
	}
}
'''

FILES["server/internal/api/recon_relay_scan.go"] = r'''package api

import (
	"encoding/json"
	"net/http"
	"strings"
	"time"

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

const relayScanTimeout = 30 * time.Second

func (h *ReconHandler) RelayScan(w http.ResponseWriter, r *http.Request) {
	var req recon.RelayScanRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid JSON", http.StatusBadRequest)
		return
	}
	req.Host = strings.TrimSpace(req.Host)
	if req.Host == "" {
		http.Error(w, "host required", http.StatusBadRequest)
		return
	}
	host, err := recon.NormalizeHost(req.Host)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	req.Host = host

	if recon.HostReachable(host) {
		report, err := recon.LocalRelayScan(req)
		if err != nil {
			writeJSON(w, map[string]interface{}{"ok": false, "error": err.Error()})
			return
		}
		writeJSON(w, map[string]interface{}{"ok": true, "report": report.ToScanReport()})
		return
	}

	if h == nil || h.wsHub == nil {
		writeJSON(w, relayScanFallback(host, "", "", "fleet websocket hub unavailable"))
		return
	}

	seedID, seedName, ok := pickSpreadSeedAgent(h.wsHub, host, "")
	if !ok {
		writeJSON(w, relayScanFallback(host, "", "", "no online fleet agent on nearby /24 to relay scan for "+host))
		return
	}

	ch := h.wsHub.AwaitCommandResult(seedID, "recon_relay_scan")
	args := map[string]interface{}{"command": host}
	if req.UDPGuess {
		args["path"] = "true"
	}
	if err := h.wsHub.SendAgentCommand(seedID, "recon_relay_scan", args); err != nil {
		h.wsHub.CancelAwait(seedID, "recon_relay_scan")
		writeJSON(w, map[string]interface{}{
			"ok": false, "agent_id": seedID, "agent_name": seedName, "error": err.Error(),
			"report": relayScanFallbackReport(host, seedName, "relay dispatch failed: "+err.Error()),
		})
		return
	}

	select {
	case payload := <-ch:
		success, _ := payload["success"].(bool)
		msg, _ := payload["message"].(string)
		if !success {
			writeJSON(w, map[string]interface{}{
				"ok": false, "agent_id": seedID, "agent_name": seedName, "error": strings.TrimSpace(msg),
				"report": relayScanFallbackReport(host, seedName, strings.TrimSpace(msg)),
			})
			return
		}
		shell := recon.RelayScanReport{
			Host: host, LocalReachable: false, ScannedVia: "relay",
			RelayAgentID: seedID, RelayAgentName: seedName,
		}
		merged, err := recon.MergeAgentRelayScan(shell, []byte(msg), req.UDPGuess)
		if err != nil {
			writeJSON(w, map[string]interface{}{
				"ok": false, "agent_id": seedID, "agent_name": seedName,
				"error": "invalid agent relay payload: " + err.Error(),
			})
			return
		}
		writeJSON(w, map[string]interface{}{
			"ok": true, "agent_id": seedID, "agent_name": seedName, "report": merged.ToScanReport(),
		})
	case <-time.After(relayScanTimeout):
		h.wsHub.CancelAwait(seedID, "recon_relay_scan")
		writeJSON(w, map[string]interface{}{
			"ok": false, "agent_id": seedID, "agent_name": seedName, "error": "relay scan timed out",
			"report": relayScanFallbackReport(host, seedName, "relay scan timed out after 30s"),
		})
	}
}

func relayScanFallback(host, agentID, agentName, message string) map[string]interface{} {
	out := map[string]interface{}{"ok": false, "error": message, "report": relayScanFallbackReport(host, agentName, message)}
	if agentID != "" {
		out["agent_id"] = agentID
	}
	if agentName != "" {
		out["agent_name"] = agentName
	}
	return out
}

func relayScanFallbackReport(host, relayVia, message string) *recon.ScanReport {
	return &recon.ScanReport{Host: host, ScannedAt: time.Now().UTC(), RelayVia: relayVia, Message: message}
}
'''

FILES["server/internal/api/recon_relay_scan_test.go"] = r'''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)
	}
}
'''

FILES["agent/client/recon_relay_scan_test.go"] = r'''package client

import (
	"encoding/json"
	"testing"

	"crypto-miner-agent/deploy"
)

func TestHandleReconRelayScanCommand(t *testing.T) {
	deploy.SetProbePortsHook(func(host string, ports []int) []int {
		if host == "10.1.2.3" {
			return []int{22, 445}
		}
		return nil
	})
	t.Cleanup(func() { deploy.SetProbePortsHook(nil) })
	deploy.SetRelayUDPGuessHook(func(host string, ports []int) []deploy.RelayUDPHint {
		return []deploy.RelayUDPHint{{Port: 51820, Open: true, Service: "wireguard"}}
	})
	t.Cleanup(func() { deploy.SetRelayUDPGuessHook(nil) })

	var gotAction string
	var gotSuccess bool
	var gotMsg string
	c := &AgentClient{}
	c.commandResultHook = func(action string, success bool, message string) {
		gotAction, gotSuccess, gotMsg = action, success, message
	}

	if !c.handleReconCommand("recon_relay_scan", "10.1.2.3", "true") {
		t.Fatal("expected handled")
	}
	if gotAction != "recon_relay_scan" || !gotSuccess {
		t.Fatalf("action=%s success=%v", gotAction, gotSuccess)
	}
	var parsed map[string]interface{}
	if err := json.Unmarshal([]byte(gotMsg), &parsed); err != nil {
		t.Fatal(err)
	}
	ports, _ := parsed["open_ports"].([]interface{})
	if len(ports) != 2 {
		t.Fatalf("ports=%v", ports)
	}
}

func TestHandleReconRelayScanMissingHost(t *testing.T) {
	c := &AgentClient{}
	c.commandResultHook = func(action string, success bool, _ string) {
		if action != "recon_relay_scan" || success {
			t.Fatalf("action=%s success=%v", action, success)
		}
	}
	if !c.handleReconCommand("recon_relay_scan", "", "") {
		t.Fatal("expected handled")
	}
}
'''


def patch_file(rel: str, old: str, new: str) -> None:
    path = ROOT / rel
    text = path.read_text(encoding="utf-8")
    if old not in text:
        if new.strip() in text:
            return
        raise SystemExit(f"patch miss in {rel}: {old[:60]!r}")
    path.write_text(text.replace(old, new, 1), encoding="utf-8", newline="\n")


def main() -> None:
    for rel, content in FILES.items():
        path = ROOT / rel
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(content, encoding="utf-8", newline="\n")
        print("wrote", rel)

    types_path = ROOT / "server/internal/recon/types.go"
    types_text = types_path.read_text(encoding="utf-8")
    if "PathTracerHints" not in types_text:
        types_text = types_text.replace(
            "\tRelayVia        string                 `json:\"relay_via,omitempty\"`\n"
            "\tRecommendations []DeployRecommendation `json:\"recommendations,omitempty\"`\n}",
            "\tRelayVia        string                 `json:\"relay_via,omitempty\"`\n"
            "\tUDPHints        []UDPHint              `json:\"udp_hints,omitempty\"`\n"
            "\tPathTracerHints []string               `json:\"path_tracer_hints,omitempty\"`\n"
            "\tMessage         string                 `json:\"message,omitempty\"`\n"
            "\tRecommendations []DeployRecommendation `json:\"recommendations,omitempty\"`\n}",
            1,
        )
        types_path.write_text(types_text, encoding="utf-8", newline="\n")

    scan_path = ROOT / "server/internal/recon/scan.go"
    if "func NormalizeHost" not in scan_path.read_text(encoding="utf-8"):
        patch_file(
            "server/internal/recon/scan.go",
            "func Scan(req ScanRequest) (*ScanReport, error) {\n\thost, err := normalizeOwnedHost(req.Host)",
            "func NormalizeHost(host string) (string, error) {\n\treturn normalizeOwnedHost(host)\n}\n\n"
            "func Scan(req ScanRequest) (*ScanReport, error) {\n\thost, err := normalizeOwnedHost(req.Host)",
        )

    router_path = ROOT / "server/internal/api/router.go"
    router_text = router_path.read_text(encoding="utf-8")
    if '"/recon/relay-scan"' not in router_text:
        patch_file(
            "server/internal/api/router.go",
            '\t\tr.Post("/recon/scan", reconHandler.Scan)\n',
            '\t\tr.Post("/recon/scan", reconHandler.Scan)\n\t\tr.Post("/recon/relay-scan", reconHandler.RelayScan)\n',
        )

    cmd_path = ROOT / "agent/client/commands_common.go"
    cmd_text = cmd_path.read_text(encoding="utf-8")
    if "recon_relay_scan" not in cmd_text:
        patch_file(
        "agent/client/commands_common.go",
        "func (c *AgentClient) handleReconCommand(action, command string) bool {\n",
        "func (c *AgentClient) handleReconCommand(action, command, path string) bool {\n"
        "\tif action == \"recon_relay_scan\" {\n"
        "\t\thost := strings.TrimSpace(command)\n"
        "\t\tif host == \"\" {\n"
        "\t\t\tc.sendCommandResult(action, false, \"host is required in command field\")\n"
        "\t\t\treturn true\n"
        "\t\t}\n"
        "\t\tudpGuess := strings.EqualFold(strings.TrimSpace(path), \"true\")\n"
        "\t\tout, err := deploy.RunRelayScanJSON(host, udpGuess)\n"
        "\t\tif err != nil {\n"
        "\t\t\tc.sendCommandResult(action, false, err.Error())\n"
        "\t\t\treturn true\n"
        "\t\t}\n"
        "\t\tc.sendCommandResult(action, true, out)\n"
        "\t\treturn true\n"
        "\t}\n",
        )

    client_path = ROOT / "agent/client/client.go"
    if "handleReconCommand(action, command, path)" not in client_path.read_text(encoding="utf-8"):
        patch_file(
            "agent/client/client.go",
            "\t\tif c.handleReconCommand(action, command) {\n",
            "\t\tif c.handleReconCommand(action, command, path) {\n",
        )

    natpunch = (ROOT / "agent/deploy/natpunch.go").read_text(encoding="utf-8")
    if "SetProbePortsHook" not in natpunch:
        patch_file(
            "agent/deploy/natpunch.go",
            "var probePortsFn func(host string, ports []int) []int\n\nfunc probePorts",
            "var probePortsFn func(host string, ports []int) []int\n\n"
            "func SetProbePortsHook(fn func(host string, ports []int) []int) { probePortsFn = fn }\n\n"
            "func probePorts",
        )

    print("patches applied")


def run(cmd: list[str], cwd: Path | None = None) -> None:
    print("+", " ".join(cmd))
    subprocess.run(cmd, cwd=cwd or ROOT, check=True)


def main_and_ship() -> None:
    main()

    paths = [
        "server/internal/recon/relay_scan.go",
        "server/internal/recon/relay_scan_test.go",
        "server/internal/recon/types.go",
        "server/internal/recon/scan.go",
        "server/internal/api/recon_relay_scan.go",
        "server/internal/api/recon_relay_scan_test.go",
        "server/internal/api/router.go",
        "agent/deploy/relay_scan.go",
        "agent/deploy/natpunch.go",
        "agent/client/commands_common.go",
        "agent/client/client.go",
        "agent/client/recon_relay_scan_test.go",
        "server/web/src/types/recon.ts",
    ]
    for rel in paths:
        if not (ROOT / rel).exists():
            raise SystemExit(f"missing before commit: {rel}")

    run(["git", "add", *paths])
    msg = (
        "Recon network batch 2: fleet relay scan and UDP hints.\n\n"
        "POST /api/v1/recon/relay-scan probes from the dashboard when reachable, "
        "otherwise dispatches recon_relay_scan to a same-/24 online agent. "
        "Agents reuse probePortsFn for TCP and optionally UDP 53/51820 with "
        "dns/wireguard Path Tracer tags."
    )
    run(["git", "commit", "-m", msg])
    run(["git", "push", "origin", "HEAD"])
    out = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT, text=True).strip()
    print("COMMIT_HASH=" + out)


if __name__ == "__main__":
    # Include agent relay_scan.go in FILES if missing above
    FILES.setdefault(
        "agent/deploy/relay_scan.go",
        (ROOT / "agent/deploy/relay_scan.go").read_text(encoding="utf-8")
        if (ROOT / "agent/deploy/relay_scan.go").exists()
        else r'''package deploy

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

var RelayScanUDPPorts = []int{53, 51820}

type RelayScanResult struct {
	Host      string         `json:"host"`
	OpenPorts []int          `json:"open_ports"`
	UDPHints  []RelayUDPHint `json:"udp_hints,omitempty"`
}

type RelayUDPHint struct {
	Port    int    `json:"port"`
	Open    bool   `json:"open"`
	Service string `json:"service"`
}

var udpGuessPortsFn func(host string, ports []int) []RelayUDPHint

func SetRelayUDPGuessHook(fn func(host string, ports []int) []RelayUDPHint) { udpGuessPortsFn = fn }

func RunRelayScan(host string, udpGuess bool) RelayScanResult {
	host = strings.TrimSpace(host)
	open := probePorts(host, SubnetReconPorts)
	result := RelayScanResult{Host: host, OpenPorts: append([]int(nil), open...)}
	if udpGuess {
		result.UDPHints = guessRelayUDPPorts(host, RelayScanUDPPorts)
	}
	return result
}

func RunRelayScanJSON(host string, udpGuess bool) (string, error) {
	b, err := json.Marshal(RunRelayScan(host, udpGuess))
	if err != nil {
		return "", err
	}
	return string(b), nil
}

func guessRelayUDPPorts(host string, ports []int) []RelayUDPHint {
	if udpGuessPortsFn != nil {
		return udpGuessPortsFn(host, ports)
	}
	var out []RelayUDPHint
	for _, port := range ports {
		out = append(out, RelayUDPHint{Port: port, Open: probeRelayUDPQuick(host, port), Service: relayUDPServiceLabel(port)})
	}
	return out
}

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

func probeRelayUDPQuick(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 := relayUDPProbePayload(port); len(b) > 0 {
		_, _ = conn.Write(b)
	}
	buf := make([]byte, 512)
	n, err := conn.Read(buf)
	return err == nil && n > 0
}

func relayUDPProbePayload(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
	}
}
''',
    )

    recon_ts = ROOT / "server/web/src/types/recon.ts"
    text = recon_ts.read_text(encoding="utf-8")
    if "relay_via?:" not in text:
        text = text.replace(
            "  crawl?: ReconCrawlReport;\n  recommendations?: ReconDeployRecommendation[];\n}",
            "  crawl?: ReconCrawlReport;\n  relay_via?: string;\n"
            "  udp_hints?: { port: number; open: boolean; service?: string }[];\n"
            "  path_tracer_hints?: string[];\n  message?: string;\n"
            "  recommendations?: ReconDeployRecommendation[];\n}",
            1,
        )
        recon_ts.write_text(text, encoding="utf-8", newline="\n")

    try:
        main_and_ship()
    except subprocess.CalledProcessError as exc:
        sys.exit(exc.returncode)