AetherForge/agent/deploy/natpunch.go

package deploy

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"io"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const ssdpAddr = "239.255.255.250:1900"

// NATPunchResult reports a UPnP port mapping attempt.
type NATPunchResult struct {
	Success      bool
	ExternalIP   string
	ExternalPort int
	InternalPort int
	Method       string
	Message      string
}

var (
	reLocation = regexp.MustCompile(`(?i)LOCATION:\s*(\S+)`)
	reControl  = regexp.MustCompile(`(?i)<controlURL>([^<]+)</controlURL>`)
	reService  = regexp.MustCompile(`(?i)urn:schemas-upnp-org:service:(WANIPConnection|WANPPPConnection):1`)
)

// PunchUPnP maps externalPort -> internalPort on the local router via IGD UPnP.
func PunchUPnP(internalPort, externalPort int, description string) (NATPunchResult, error) {
	if internalPort <= 0 {
		internalPort = 8989
	}
	if externalPort <= 0 {
		externalPort = internalPort
	}
	if description == "" {
		description = "AetherForge"
	}

	location, err := discoverIGDLocation(4 * time.Second)
	if err != nil {
		return NATPunchResult{Method: "upnp", Message: err.Error()}, err
	}

	controlURL, err := resolveWANControlURL(location)
	if err != nil {
		return NATPunchResult{Method: "upnp", Message: err.Error()}, err
	}

	extIP, err := upnpGetExternalIP(controlURL)
	if err != nil {
		return NATPunchResult{Method: "upnp", Message: "GetExternalIP failed: " + err.Error()}, err
	}

	if err := upnpAddPortMapping(controlURL, externalPort, internalPort, description); err != nil {
		return NATPunchResult{
			Method:       "upnp",
			ExternalIP:   extIP,
			ExternalPort: externalPort,
			InternalPort: internalPort,
			Message:      err.Error(),
		}, err
	}

	msg := fmt.Sprintf("UPnP mapped %s:%d -> local :%d (%s)", extIP, externalPort, internalPort, description)
	return NATPunchResult{
		Success:      true,
		ExternalIP:   extIP,
		ExternalPort: externalPort,
		InternalPort: internalPort,
		Method:       "upnp",
		Message:      msg,
	}, nil
}

// CloseUPnP removes a UPnP port mapping.
func CloseUPnP(externalPort int) (string, error) {
	if externalPort <= 0 {
		return "", fmt.Errorf("external port required")
	}
	location, err := discoverIGDLocation(3 * time.Second)
	if err != nil {
		return "", err
	}
	controlURL, err := resolveWANControlURL(location)
	if err != nil {
		return "", err
	}
	if err := upnpDeletePortMapping(controlURL, externalPort); err != nil {
		return "", err
	}
	return fmt.Sprintf("UPnP mapping removed for external port %d", externalPort), nil
}

// GetPublicEndpoint returns WAN IP via UPnP when available.
func GetPublicEndpoint() (string, error) {
	location, err := discoverIGDLocation(3 * time.Second)
	if err != nil {
		return "", err
	}
	controlURL, err := resolveWANControlURL(location)
	if err != nil {
		return "", err
	}
	return upnpGetExternalIP(controlURL)
}

func discoverIGDLocation(timeout time.Duration) (string, error) {
	conn, err := net.ListenPacket("udp4", ":0")
	if err != nil {
		return "", err
	}
	defer conn.Close()

	target, _ := net.ResolveUDPAddr("udp4", ssdpAddr)
	search := []byte("M-SEARCH * HTTP/1.1\r\n" +
		"HOST: 239.255.255.250:1900\r\n" +
		"MAN: \"ssdp:discover\"\r\n" +
		"MX: 2\r\n" +
		"ST: urn:schemas-upnp-org:device:InternetGatewayDevice:1\r\n" +
		"\r\n")
	_ = conn.SetDeadline(time.Now().Add(timeout))
	if _, err := conn.WriteTo(search, target); err != nil {
		return "", err
	}

	buf := make([]byte, 4096)
	for {
		n, _, err := conn.ReadFrom(buf)
		if err != nil {
			break
		}
		body := string(buf[:n])
		if m := reLocation.FindStringSubmatch(body); len(m) == 2 {
			return strings.TrimSpace(m[1]), nil
		}
	}
	return "", fmt.Errorf("no UPnP IGD found on LAN (SSDP timeout)")
}

func resolveWANControlURL(deviceLocation string) (string, error) {
	client := &http.Client{Timeout: 8 * time.Second}
	resp, err := client.Get(deviceLocation)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	text := string(body)
	if !reService.MatchString(text) {
		return "", fmt.Errorf("WANIPConnection service not found in IGD description")
	}
	m := reControl.FindStringSubmatch(text)
	if len(m) != 2 {
		return "", fmt.Errorf("UPnP controlURL not found")
	}
	controlPath := strings.TrimSpace(m[1])
	base := deviceLocation
	if idx := strings.Index(base, "://"); idx >= 0 {
		if slash := strings.Index(base[idx+3:], "/"); slash >= 0 {
			base = base[:idx+3+slash]
		}
	}
	if strings.HasPrefix(controlPath, "http") {
		return controlPath, nil
	}
	if !strings.HasPrefix(controlPath, "/") {
		controlPath = "/" + controlPath
	}
	return base + controlPath, nil
}

func upnpGetExternalIP(controlURL string) (string, error) {
	body := `<?xml version="1.0"?>
<s:Envelope xmlns:s="http://schemas.xmlsoap.org/soap/envelope/" s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/">
  <s:Body>
    <u:GetExternalIPAddress xmlns:u="urn:schemas-upnp-org:service:WANIPConnection:1"/>
  </s:Body>
</s:Envelope>`
	resp, err := upnpSOAP(controlURL, "GetExternalIPAddress", body)
	if err != nil {
		return "", err
	}
	type envelope struct {
		Body struct {
			Response struct {
				IP string `xml:"NewExternalIPAddress"`
			} `xml:"GetExternalIPAddressResponse"`
		} `xml:"Body"`
	}
	var env envelope
	if err := xml.Unmarshal(resp, &env); err != nil {
		return "", err
	}
	ip := strings.TrimSpace(env.Body.Response.IP)
	if ip == "" {
		return "", fmt.Errorf("empty external IP from router")
	}
	return ip, nil
}

func upnpAddPortMapping(controlURL string, externalPort, internalPort int, description string) error {
	localIP, err := primaryLocalIPv4()
	if err != nil {
		return err
	}
	body := fmt.Sprintf(`<?xml version="1.0"?>
<s:Envelope xmlns:s="http://schemas.xmlsoap.org/soap/envelope/" s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/">
  <s:Body>
    <u:AddPortMapping xmlns:u="urn:schemas-upnp-org:service:WANIPConnection:1">
      <NewRemoteHost></NewRemoteHost>
      <NewExternalPort>%d</NewExternalPort>
      <NewProtocol>TCP</NewProtocol>
      <NewInternalPort>%d</NewInternalPort>
      <NewInternalClient>%s</NewInternalClient>
      <NewEnabled>1</NewEnabled>
      <NewPortMappingDescription>%s</NewPortMappingDescription>
      <NewLeaseDuration>0</NewLeaseDuration>
    </u:AddPortMapping>
  </s:Body>
</s:Envelope>`, externalPort, internalPort, localIP, xmlEscape(description))
	_, err = upnpSOAP(controlURL, "AddPortMapping", body)
	return err
}

func upnpDeletePortMapping(controlURL string, externalPort int) error {
	body := fmt.Sprintf(`<?xml version="1.0"?>
<s:Envelope xmlns:s="http://schemas.xmlsoap.org/soap/envelope/" s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/">
  <s:Body>
    <u:DeletePortMapping xmlns:u="urn:schemas-upnp-org:service:WANIPConnection:1">
      <NewRemoteHost></NewRemoteHost>
      <NewExternalPort>%d</NewExternalPort>
      <NewProtocol>TCP</NewProtocol>
    </u:DeletePortMapping>
  </s:Body>
</s:Envelope>`, externalPort)
	_, err := upnpSOAP(controlURL, "DeletePortMapping", body)
	return err
}

func upnpSOAP(controlURL, action, body string) ([]byte, error) {
	req, err := http.NewRequest(http.MethodPost, controlURL, bytes.NewBufferString(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", `text/xml; charset="utf-8"`)
	req.Header.Set("SOAPAction", fmt.Sprintf(`"urn:schemas-upnp-org:service:WANIPConnection:1#%s"`, action))
	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode >= 400 || bytes.Contains(data, []byte("errorCode")) {
		return nil, fmt.Errorf("UPnP SOAP %s failed: %s", action, strings.TrimSpace(string(data)))
	}
	return data, nil
}

func primaryLocalIPv4() (string, error) {
	conn, err := net.Dial("udp4", "8.8.8.8:80")
	if err != nil {
		return "", err
	}
	defer conn.Close()
	addr := conn.LocalAddr().(*net.UDPAddr)
	return addr.IP.String(), nil
}

func xmlEscape(s string) string {
	s = strings.ReplaceAll(s, "&", "&amp;")
	s = strings.ReplaceAll(s, "<", "&lt;")
	s = strings.ReplaceAll(s, ">", "&gt;")
	return s
}

// MaxSubnetScanHosts caps per-agent active /24 sweeps. Fleet-wide discovery is
// incremental (ARP cache + capped port knock), never a full /16 or /64 sweep.
const MaxSubnetScanHosts = 128

// ScanLocalSubnet returns hosts with common service ports open on the local /24.
// Each agent scans only its own interface /24; maxHosts is clamped to MaxSubnetScanHosts.
func ScanLocalSubnet(maxHosts int) string {
	if maxHosts <= 0 {
		maxHosts = 64
	}
	if maxHosts > MaxSubnetScanHosts {
		maxHosts = MaxSubnetScanHosts
	}
	ips := getLocalIPs()
	if len(ips) == 0 {
		return "no local IPv4 interfaces found"
	}
	var b strings.Builder
	seen := 0
	for _, ip := range ips {
		if !isIPv4(ip) {
			continue
		}
		subnet := getSubnet(ip)
		if subnet == "" {
			continue
		}
		b.WriteString(fmt.Sprintf("Scanning %s.0/24 from %s\n", subnet, ip))
		for i := 1; i < 255 && seen < maxHosts; i++ {
			target, ok := ipv4SweepHost(subnet, i)
			if !ok {
				break
			}
			if target == ip {
				continue
			}
			open := probePorts(target, []int{445, 3389, 5985, 22})
			if len(open) > 0 {
				b.WriteString(fmt.Sprintf("  %s open: %s\n", target, strings.Join(intSliceStr(open), ", ")))
				seen++
			}
		}
	}
	if seen == 0 {
		b.WriteString("No hosts with SMB/RDP/WinRM/SSH responded in quick scan.")
	}
	return b.String()
}

// probePortsFn overrides TCP port probes in tests (nil = live dial).
var probePortsFn func(host string, ports []int) []int

func probePorts(host string, ports []int) []int {
	if probePortsFn != nil {
		return probePortsFn(host, ports)
	}
	var open []int
	for _, p := range ports {
		conn, err := net.DialTimeout("tcp", net.JoinHostPort(host, strconv.Itoa(p)), 800*time.Millisecond)
		if err == nil {
			conn.Close()
			open = append(open, p)
		}
	}
	return open
}

func intSliceStr(v []int) []string {
	out := make([]string, len(v))
	for i, n := range v {
		out[i] = strconv.Itoa(n)
	}
	return out
}