Pulse/pkg/discovery/discovery.go

package discovery

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/rs/zerolog/log"
)

// DiscoveredServer represents a discovered Proxmox/PBS server
type DiscoveredServer struct {
	IP       string `json:"ip"`
	Port     int    `json:"port"`
	Type     string `json:"type"` // "pve" or "pbs"
	Version  string `json:"version"`
	Hostname string `json:"hostname,omitempty"`
	Release  string `json:"release,omitempty"`
}

// DiscoveryResult contains all discovered servers
type DiscoveryResult struct {
	Servers []DiscoveredServer `json:"servers"`
	Errors  []string           `json:"errors,omitempty"`
}

// Scanner handles network scanning for Proxmox/PBS servers
type Scanner struct {
	timeout     time.Duration
	concurrent  int
	httpClient  *http.Client
}

// NewScanner creates a new network scanner
func NewScanner() *Scanner {
	return &Scanner{
		timeout:    1 * time.Second,  // Reduced timeout for faster scanning
		concurrent: 50, // Increased concurrent workers for faster scanning
		httpClient: &http.Client{
			Timeout: 2 * time.Second,  // Reduced HTTP timeout
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
				MaxIdleConns:    100,
				MaxConnsPerHost: 10,
			},
		},
	}
}

// DiscoverServers scans the network for Proxmox VE and PBS servers
func (s *Scanner) DiscoverServers(ctx context.Context, subnet string) (*DiscoveryResult, error) {
	log.Info().Str("subnet", subnet).Msg("Starting network discovery")
	
	// Parse subnet
	var ipNet *net.IPNet
	
	if subnet == "" || subnet == "auto" {
		// Auto-detect local subnet
		ipNet = s.getLocalSubnet()
		if ipNet == nil {
			return nil, fmt.Errorf("failed to auto-detect local subnet")
		}
		log.Info().Str("detected", ipNet.String()).Msg("Auto-detected local subnet")
	} else {
		// Parse provided subnet
		_, parsedNet, err := net.ParseCIDR(subnet)
		if err != nil {
			return nil, fmt.Errorf("invalid subnet: %w", err)
		}
		ipNet = parsedNet
	}
	
	// Check subnet size - limit to /24 or smaller for safety
	ones, bits := ipNet.Mask.Size()
	if ones < 24 && bits == 32 { // IPv4 with more than 256 addresses
		log.Warn().Str("subnet", ipNet.String()).Msg("Subnet too large, limiting to /24")
		// Convert to /24
		ipNet.Mask = net.CIDRMask(24, 32)
	}

	// Generate list of IPs to scan
	ips := s.generateIPs(ipNet)
	log.Info().Int("count", len(ips)).Msg("IPs to scan")

	// Create channels for work distribution
	ipChan := make(chan string, len(ips))
	resultChan := make(chan *DiscoveredServer, len(ips))
	errorChan := make(chan string, len(ips))

	// Start workers
	var wg sync.WaitGroup
	for i := 0; i < s.concurrent; i++ {
		wg.Add(1)
		go s.scanWorker(ctx, &wg, ipChan, resultChan, errorChan)
	}

	// Send IPs to scan
	for _, ip := range ips {
		ipChan <- ip
	}
	close(ipChan)

	// Wait for workers to finish
	go func() {
		wg.Wait()
		close(resultChan)
		close(errorChan)
	}()

	// Collect results
	result := &DiscoveryResult{
		Servers: []DiscoveredServer{},
		Errors:  []string{},
	}

	done := false
	for !done {
		select {
		case server, ok := <-resultChan:
			if !ok {
				done = true
				break
			}
			if server != nil {
				result.Servers = append(result.Servers, *server)
			}
		case errMsg, ok := <-errorChan:
			if ok && errMsg != "" {
				result.Errors = append(result.Errors, errMsg)
			}
		case <-ctx.Done():
			return result, ctx.Err()
		}
	}

	log.Info().
		Int("found", len(result.Servers)).
		Int("errors", len(result.Errors)).
		Msg("Discovery completed")

	return result, nil
}

// scanWorker scans IPs from the channel
func (s *Scanner) scanWorker(ctx context.Context, wg *sync.WaitGroup, ipChan <-chan string, resultChan chan<- *DiscoveredServer, errorChan chan<- string) {
	defer wg.Done()

	for ip := range ipChan {
		select {
		case <-ctx.Done():
			return
		default:
			// Check Proxmox VE (port 8006)
			if server := s.checkServer(ctx, ip, 8006, "pve"); server != nil {
				resultChan <- server
			}

			// Check PBS (port 8007)
			if server := s.checkServer(ctx, ip, 8007, "pbs"); server != nil {
				resultChan <- server
			}
		}
	}
}

// checkServer checks if a server is running at the given IP and port
func (s *Scanner) checkServer(ctx context.Context, ip string, port int, serverType string) *DiscoveredServer {
	// First check if port is open
	address := fmt.Sprintf("%s:%d", ip, port)
	conn, err := net.DialTimeout("tcp", address, s.timeout)
	if err != nil {
		return nil // Port not open
	}
	conn.Close()

	// Port is open - this is likely a Proxmox/PBS server
	// Since most installations require auth for version endpoint,
	// we'll return it as a discovered server based on the port alone
	
	log.Info().
		Str("ip", ip).
		Int("port", port).
		Str("type", serverType).
		Msg("Found potential server (port open)")

	server := &DiscoveredServer{
		IP:      ip,
		Port:    port,
		Type:    serverType,
		Version: "Unknown", // Will be determined after auth
	}

	// Try to get version without auth (some installations allow it)
	url := fmt.Sprintf("https://%s:%d/api2/json/version", ip, port)
	
	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err == nil {
		resp, err := s.httpClient.Do(req)
		if err == nil {
			defer resp.Body.Close()
			
			// Only try to parse if we got a successful response
			if resp.StatusCode == 200 {
				var versionResp struct {
					Data struct {
						Version string `json:"version"`
						Release string `json:"release,omitempty"`
					} `json:"data"`
				}
				
				if err := json.NewDecoder(resp.Body).Decode(&versionResp); err == nil && versionResp.Data.Version != "" {
					server.Version = versionResp.Data.Version
					server.Release = versionResp.Data.Release
					
					log.Info().
						Str("ip", ip).
						Int("port", port).
						Str("version", server.Version).
						Msg("Got server version without auth")
				}
			}
		}
	}

	// Try to resolve hostname via reverse DNS
	names, err := net.LookupAddr(ip)
	if err == nil && len(names) > 0 {
		// Use the first hostname, remove trailing dot if present
		hostname := strings.TrimSuffix(names[0], ".")
		server.Hostname = hostname
		log.Debug().Str("ip", ip).Str("hostname", hostname).Msg("Resolved hostname via DNS")
	}

	return server
}

// getProxmoxHostname tries to get the hostname of a Proxmox VE server
func (s *Scanner) getProxmoxHostname(ctx context.Context, ip string, port int) string {
	url := fmt.Sprintf("https://%s:%d/api2/json/nodes", ip, port)
	
	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return ""
	}

	resp, err := s.httpClient.Do(req)
	if err != nil {
		return ""
	}
	defer resp.Body.Close()

	var nodesResp struct {
		Data []struct {
			Node string `json:"node"`
		} `json:"data"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&nodesResp); err != nil {
		return ""
	}

	if len(nodesResp.Data) > 0 {
		return nodesResp.Data[0].Node
	}

	return ""
}

// getPBSHostname tries to get the hostname of a PBS server
func (s *Scanner) getPBSHostname(ctx context.Context, ip string, port int) string {
	url := fmt.Sprintf("https://%s:%d/api2/json/nodes", ip, port)
	
	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return ""
	}

	resp, err := s.httpClient.Do(req)
	if err != nil {
		return ""
	}
	defer resp.Body.Close()

	var nodesResp struct {
		Data []struct {
			Node string `json:"node"`
		} `json:"data"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&nodesResp); err != nil {
		return ""
	}

	if len(nodesResp.Data) > 0 {
		return nodesResp.Data[0].Node
	}

	return ""
}

// generateIPs generates all IPs in a subnet
func (s *Scanner) generateIPs(ipNet *net.IPNet) []string {
	var ips []string
	
	// Get the starting IP
	ip := ipNet.IP.Mask(ipNet.Mask)
	
	// Calculate the number of hosts
	ones, bits := ipNet.Mask.Size()
	hostBits := bits - ones
	numHosts := 1 << hostBits
	
	// Skip network and broadcast addresses for common subnets
	start := 1
	end := numHosts - 1
	if numHosts > 256 {
		// For larger subnets, scan everything
		start = 0
		end = numHosts
	}
	
	// Limit to maximum 1024 IPs to avoid scanning huge networks
	if end-start > 1024 {
		end = start + 1024
		log.Warn().Int("limited_to", 1024).Msg("Limiting scan to first 1024 IPs")
	}
	
	for i := start; i < end; i++ {
		// Calculate IP
		currIP := make(net.IP, len(ip))
		copy(currIP, ip)
		
		// Add offset to IP address
		offset := i
		for j := len(currIP) - 1; j >= 0 && offset > 0; j-- {
			currIP[j] += byte(offset & 0xFF)
			offset >>= 8
		}
		
		// Skip common non-server IPs
		lastOctet := currIP[len(currIP)-1]
		if lastOctet == 0 || lastOctet == 255 {
			continue // Skip network and broadcast
		}
		
		ips = append(ips, currIP.String())
	}
	
	return ips
}

// getLocalSubnet attempts to detect the local subnet
func (s *Scanner) getLocalSubnet() *net.IPNet {
	interfaces, err := net.Interfaces()
	if err != nil {
		return nil
	}

	for _, iface := range interfaces {
		// Skip loopback and down interfaces
		if iface.Flags&net.FlagLoopback != 0 || iface.Flags&net.FlagUp == 0 {
			continue
		}

		addrs, err := iface.Addrs()
		if err != nil {
			continue
		}

		for _, addr := range addrs {
			if ipNet, ok := addr.(*net.IPNet); ok && ipNet.IP.To4() != nil {
				// Found an IPv4 address
				if !ipNet.IP.IsLoopback() && !ipNet.IP.IsLinkLocalUnicast() {
					// Convert to /24 subnet for auto-detection
					// This ensures we scan a reasonable range
					ip := ipNet.IP.To4()
					if ip != nil {
						// Create a /24 subnet from the IP
						ip[3] = 0 // Set last octet to 0
						return &net.IPNet{
							IP:   ip,
							Mask: net.CIDRMask(24, 32),
						}
					}
				}
			}
		}
	}

	// Default to common subnet if detection fails
	_, defaultNet, _ := net.ParseCIDR("192.168.1.0/24")
	return defaultNet
}