Improve discovery classification heuristics

Refs #551

frontend-modern/src/components/Settings/Settings.tsx

@@ -457,24 +457,45 @@ const Settings: Component<SettingsProps> = (props) => {
       }
     });
 
+    const recognizedTypes = ['pve', 'pbs', 'pmg'] as const;
+    type RecognizedType = (typeof recognizedTypes)[number];
+    const isRecognizedType = (value: string): value is RecognizedType =>
+      (recognizedTypes as readonly string[]).includes(value);
+
     const normalized = servers
       .map((server): DiscoveredServer | null => {
         const ip = (server.ip || '').trim();
-        const type = (server.type || '').toLowerCase();
+        let type = (server.type || '').toLowerCase();
-        const port = typeof server.port === 'number' ? server.port : type === 'pbs' ? 8007 : 8006;
-
-        if (!ip || (type !== 'pve' && type !== 'pbs')) {
-          return null;
-        }
-
         const hostname = (server.hostname || server.name || '').trim();
         const version = (server.version || '').trim();
         const release = (server.release || '').trim();
 
+        if (!isRecognizedType(type)) {
+          const metadata = `${hostname} ${version} ${release}`.toLowerCase();
+          if (metadata.includes('pmg') || metadata.includes('mail gateway')) {
+            type = 'pmg';
+          } else if (metadata.includes('pbs') || metadata.includes('backup server')) {
+            type = 'pbs';
+          } else if (metadata.includes('pve') || metadata.includes('virtual environment')) {
+            type = 'pve';
+          }
+        }
+
+        if (!ip || !isRecognizedType(type)) {
+          return null;
+        }
+
+        const port =
+          typeof server.port === 'number'
+            ? server.port
+            : type === 'pbs'
+            ? 8007
+            : 8006;
+
         return {
           ip,
           port,
-          type: type as 'pve' | 'pbs',
+          type,
           version: version || 'Unknown',
           hostname: hostname || undefined,
           release: release || undefined,

pkg/discovery/discovery.go

@@ -206,21 +206,85 @@ func (s *Scanner) scanWorker(ctx context.Context, wg *sync.WaitGroup, ipChan <-c
 
 // checkPort8006 checks if port 8006 is running PMG or PVE
 func (s *Scanner) checkPort8006(ctx context.Context, ip string) *DiscoveredServer {
-	// First check if port is open
 	address := net.JoinHostPort(ip, "8006")
-	conn, err := net.DialTimeout("tcp", address, s.timeout)
-	if err != nil {
-		return nil // Port not open
-	}
-	conn.Close()
 
-	// Port is open - now detect if it's PMG or PVE
+	// First attempt a TLS handshake so we can inspect certificate metadata.
-	// PMG has statistics/mail endpoint that PVE doesn't have
+	var tlsState *tls.ConnectionState
-	// Even if auth is required, PMG will return 401, PVE will return 404
+	dialer := &net.Dialer{Timeout: s.timeout}
-	isPMG := s.isPMGServer(ctx, address)
+	tlsConn, tlsErr := tls.DialWithDialer(dialer, "tcp", address, &tls.Config{InsecureSkipVerify: true})
+	if tlsErr != nil {
+		// Fallback to a simple TCP dial to confirm the port is open.
+		conn, err := net.DialTimeout("tcp", address, s.timeout)
+		if err != nil {
+			return nil // Port not open
+		}
+		conn.Close()
+	} else {
+		state := tlsConn.ConnectionState()
+		tlsState = &state
+		tlsConn.Close()
+	}
 
 	serverType := "pve"
-	if isPMG {
+	if tlsState != nil {
+		if guess := inferTypeFromCertificate(*tlsState); guess != "" {
+			serverType = guess
+		}
+	}
+
+	version := "Unknown"
+	var release string
+
+	// Try to get version without auth (some installations allow it)
+	versionURL := fmt.Sprintf("https://%s/api2/json/version", address)
+	if req, err := http.NewRequestWithContext(ctx, "GET", versionURL, nil); err == nil {
+		if resp, err := s.httpClient.Do(req); err == nil {
+			defer resp.Body.Close()
+
+			switch resp.StatusCode {
+			case http.StatusOK:
+				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 != "" {
+					version = versionResp.Data.Version
+					release = versionResp.Data.Release
+
+					if guess := inferTypeFromMetadata(
+						versionResp.Data.Version,
+						versionResp.Data.Release,
+						resp.Header.Get("Server"),
+						resp.Header.Get("Proxmox-Product"),
+						resp.Header.Get("WWW-Authenticate"),
+						strings.Join(resp.Header.Values("Set-Cookie"), " "),
+					); guess != "" {
+						serverType = guess
+					}
+
+					log.Info().
+						Str("ip", ip).
+						Int("port", 8006).
+						Str("version", version).
+						Msg("Got server version without auth")
+				}
+			case http.StatusUnauthorized, http.StatusForbidden:
+				if guess := inferTypeFromMetadata(
+					resp.Header.Get("WWW-Authenticate"),
+					resp.Header.Get("Server"),
+					resp.Header.Get("Proxmox-Product"),
+				); guess != "" {
+					serverType = guess
+				}
+			}
+		}
+	}
+
+	// Fallback: probe PMG-specific endpoints if we still think this is a PVE server.
+	if serverType != "pmg" && s.isPMGServer(ctx, address) {
 		serverType = "pmg"
 	}
 
@@ -234,39 +298,8 @@ func (s *Scanner) checkPort8006(ctx context.Context, ip string) *DiscoveredServe
 		IP:      ip,
 		Port:    8006,
 		Type:    serverType,
-		Version: "Unknown", // Will be determined after auth
+		Version: version,
-	}
+		Release: release,
-
-	// Try to get version without auth (some installations allow it)
-	url := fmt.Sprintf("https://%s/api2/json/version", address)
-
-	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", 8006).
-						Str("version", server.Version).
-						Msg("Got server version without auth")
-				}
-			}
-		}
 	}
 
 	// Try to resolve hostname via reverse DNS
@@ -283,34 +316,116 @@ func (s *Scanner) checkPort8006(ctx context.Context, ip string) *DiscoveredServe
 
 // isPMGServer checks if a server is PMG by checking for PMG-specific endpoints
 func (s *Scanner) isPMGServer(ctx context.Context, address string) bool {
-	// Try PMG-specific endpoint: /api2/json/statistics/mail
+	endpoints := []string{
-	// PMG will return 401 (unauthorized) or 200
+		"api2/json/statistics/mail",
-	// PVE will return 404 (not found)
+		"api2/json/mail/queue",
-	url := fmt.Sprintf("https://%s/api2/json/statistics/mail", address)
+		"api2/json/mail/quarantine",
+	}
+
+	for _, endpoint := range endpoints {
+		product := s.detectProductFromEndpoint(ctx, address, endpoint)
+		if product == "pmg" {
+			log.Debug().
+				Str("address", address).
+				Str("endpoint", endpoint).
+				Msg("PMG-specific endpoint confirmed")
+			return true
+		}
+	}
+
+	return false
+}
+
+// detectProductFromEndpoint inspects an HTTP endpoint and tries to infer the product type.
+func (s *Scanner) detectProductFromEndpoint(ctx context.Context, address, endpoint string) string {
+	url := fmt.Sprintf("https://%s/%s", address, endpoint)
 
 	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
 	if err != nil {
-		return false
+		return ""
 	}
 
 	resp, err := s.httpClient.Do(req)
 	if err != nil {
-		return false
+		return ""
 	}
 	defer resp.Body.Close()
 
-	// If we get anything other than 404, it's likely PMG
+	headerProduct := inferTypeFromMetadata(
-	// PMG will return 401 (needs auth) or 200 (if auth not required)
+		resp.Header.Get("Server"),
-	// PVE will return 404 (endpoint doesn't exist)
+		resp.Header.Get("Proxmox-Product"),
-	if resp.StatusCode != 404 {
+		resp.Header.Get("WWW-Authenticate"),
-		log.Debug().
+		strings.Join(resp.Header.Values("Set-Cookie"), " "),
-			Str("address", address).
+	)
-			Int("status", resp.StatusCode).
+	if headerProduct != "" {
-			Msg("PMG-specific endpoint found - identified as PMG")
+		return headerProduct
-		return true
 	}
 
-	return false
+	// If the endpoint responded (not 404) and the path is PMG-specific, treat it as PMG.
+	if resp.StatusCode != http.StatusNotFound && strings.Contains(endpoint, "mail") {
+		return "pmg"
+	}
+
+	return ""
+}
+
+// inferTypeFromCertificate tries to determine the product based on TLS certificate metadata.
+func inferTypeFromCertificate(state tls.ConnectionState) string {
+	if len(state.PeerCertificates) == 0 {
+		return ""
+	}
+
+	cert := state.PeerCertificates[0]
+	parts := []string{cert.Subject.CommonName}
+	parts = append(parts, cert.Subject.Organization...)
+	parts = append(parts, cert.Subject.OrganizationalUnit...)
+
+	return inferTypeFromMetadata(parts...)
+}
+
+// inferTypeFromMetadata inspects textual metadata and returns a best-effort product type.
+func inferTypeFromMetadata(parts ...string) string {
+	var builder strings.Builder
+
+	for _, part := range parts {
+		part = strings.TrimSpace(part)
+		if part == "" {
+			continue
+		}
+		if builder.Len() > 0 {
+			builder.WriteByte(' ')
+		}
+		builder.WriteString(strings.ToLower(part))
+	}
+
+	combined := builder.String()
+	if combined == "" {
+		return ""
+	}
+
+	compact := strings.ReplaceAll(combined, " ", "")
+
+	switch {
+	case strings.Contains(combined, "pmg"),
+		strings.Contains(combined, "mail gateway"),
+		strings.Contains(combined, "pmgauth"),
+		strings.Contains(combined, "pmgauthcookie"),
+		strings.Contains(compact, "mailgateway"),
+		strings.Contains(compact, "pmg-api"):
+		return "pmg"
+	case strings.Contains(combined, "pbs"),
+		strings.Contains(combined, "backup server"),
+		strings.Contains(combined, "pbsauth"),
+		strings.Contains(compact, "pbs-api"):
+		return "pbs"
+	case strings.Contains(combined, "pve"),
+		strings.Contains(combined, "virtual environment"),
+		strings.Contains(combined, "pveauth"),
+		strings.Contains(compact, "pve-api"):
+		return "pve"
+	default:
+		return ""
+	}
 }
 
 // checkServer checks if a server is running at the given IP and port

pkg/discovery/discovery_test.go

@@ -0,0 +1,48 @@
+package discovery
+
+import "testing"
+
+func TestInferTypeFromMetadata(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name  string
+		parts []string
+		want  string
+	}{
+		{
+			name:  "detects PMG from auth header",
+			parts: []string{`PMGAuth realm="Proxmox Mail Gateway"`, "pmgproxy/4.0"},
+			want:  "pmg",
+		},
+		{
+			name:  "detects PVE from realm string",
+			parts: []string{`PVEAuth realm="Proxmox Virtual Environment"`, "pve-api-daemon/3.0"},
+			want:  "pve",
+		},
+		{
+			name:  "detects PBS from cookie",
+			parts: []string{"PBS", "PBSCookie=abc123", `PBSAuth realm="Proxmox Backup Server"`},
+			want:  "pbs",
+		},
+		{
+			name:  "returns empty when no markers",
+			parts: []string{"Custom Certificate", "Example Corp"},
+			want:  "",
+		},
+		{
+			name:  "tolerates compact strings",
+			parts: []string{"ProxmoxMailGateway"},
+			want:  "pmg",
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			if got := inferTypeFromMetadata(tc.parts...); got != tc.want {
+				t.Fatalf("inferTypeFromMetadata(%v) = %q, want %q", tc.parts, got, tc.want)
+			}
+		})
+	}
+}