Enhance VM entity ID generation to ensure uniqueness and handle collisions in Unraid switch integration

Author: ruaan-deysel

custom_components/unraid/api/vm_operations.py

@@ -81,17 +81,18 @@ async def get_vms(self) -> List[Dict[str, Any]]:
 
             # Collect VM information in a single command
             try:
-                # Use a single command to collect all VM information
+                # Use a more robust command that properly handles VM names with spaces and special characters
+                # We use a unique delimiter (§§§) that's unlikely to appear in VM names or XML data
                 cmd = (
                     "if [ -x /etc/rc.d/rc.libvirt ] && /etc/rc.d/rc.libvirt status | grep -q 'is currently running'; then "
-                    "  for vm in $(virsh list --all --name); do "
-                    "    if [ -n \"$vm\" ]; then "
+                    "  virsh list --all --name | while IFS= read -r vm; do "
+                    "    if [ -n \"$vm\" ] && [ \"$vm\" != \" \" ]; then "
                     "      state=$(virsh domstate \"$vm\" 2>/dev/null || echo 'unknown'); "
                     "      info=$(virsh dominfo \"$vm\" 2>/dev/null); "
-                    "      cpus=$(echo \"$info\" | grep 'CPU(s)' | awk '{print $2}'); "
-                    "      mem=$(echo \"$info\" | grep 'Max memory' | sed 's/Max memory://g' | xargs); "
-                    "      xml=$(virsh dumpxml \"$vm\" 2>/dev/null | grep -A5 \"<os>\"); "
-                    "      echo \"$vm|$state|$cpus|$mem|$xml\"; "
+                    "      cpus=$(echo \"$info\" | grep 'CPU(s)' | awk '{print $2}' | head -1); "
+                    "      mem=$(echo \"$info\" | grep 'Max memory' | sed 's/Max memory://g' | xargs | head -1); "
+                    "      xml=$(virsh dumpxml \"$vm\" 2>/dev/null | grep -A5 \"<os>\" | tr '\\n' ' '); "
+                    "      echo \"$vm§§§$state§§§$cpus§§§$mem§§§$xml\"; "
                     "    fi; "
                     "  done; "
                     "else "
@@ -107,18 +108,26 @@ async def get_vms(self) -> List[Dict[str, Any]]:
 
                 vms = []
                 for line in result.stdout.splitlines():
-                    if not line.strip() or '|' not in line:
+                    if not line.strip() or '§§§' not in line:
                         continue
 
                     try:
-                        parts = line.split('|')
+                        # Split on our unique delimiter
+                        parts = line.split('§§§')
                         if len(parts) >= 5:
                             vm_name = parts[0].strip()
                             status = VMState.parse(parts[1].strip())
-                            cpus = parts[2].strip()
-                            memory = parts[3].strip()
+                            cpus = parts[2].strip() or '0'
+                            memory = parts[3].strip() or '0'
                             xml_data = parts[4].strip()
 
+                            # Skip empty VM names (shouldn't happen but be safe)
+                            if not vm_name:
+                                _LOGGER.debug("Skipping VM with empty name")
+                                continue
+
+                            _LOGGER.debug("Processing VM: '%s' with status: '%s'", vm_name, status)
+
                             # Determine OS type from XML data
                             os_type = 'unknown'
                             xml_lower = xml_data.lower()
@@ -149,6 +158,7 @@ async def get_vms(self) -> List[Dict[str, Any]]:
                         _LOGGER.debug("Error processing VM line '%s': %s", line, str(vm_err))
                         continue
 
+                _LOGGER.debug("Successfully processed %d VMs", len(vms))
                 return vms
 
             except Exception as virsh_err:
@@ -170,23 +180,27 @@ async def _get_vms_original(self) -> List[Dict[str, Any]]:
 
             vms = []
             for line in result.stdout.splitlines():
-                if not line.strip():
+                vm_name = line.strip()
+                if not vm_name:
                     continue
 
                 try:
-                    vm_name = line.strip()
+                    _LOGGER.debug("Processing VM (fallback): '%s'", vm_name)
                     status = await self.get_vm_status(vm_name)
                     os_type = await self.get_vm_os_info(vm_name)
 
                     vms.append({
                         "name": vm_name,
                         "status": status,
-                        "os_type": os_type
+                        "os_type": os_type,
+                        "cpus": "0",  # Add default values for consistency
+                        "memory": "0"
                     })
                 except Exception as vm_err:
-                    _LOGGER.debug("Error processing VM '%s': %s", line.strip(), str(vm_err))
+                    _LOGGER.debug("Error processing VM '%s': %s", vm_name, str(vm_err))
                     continue
 
+            _LOGGER.debug("Successfully processed %d VMs (fallback)", len(vms))
             return vms
         except Exception as err:
             _LOGGER.debug("Error in original VM method: %s", str(err))
@@ -236,7 +250,9 @@ async def get_vm_os_info(self, vm_name: str) -> str:
     async def get_vm_status(self, vm_name: str) -> str:
         """Get detailed status of a specific virtual machine."""
         try:
-            result = await self.execute_command(f'virsh domstate "{vm_name}"')
+            # Use shlex.quote to properly escape VM names with special characters
+            escaped_name = shlex.quote(vm_name)
+            result = await self.execute_command(f'virsh domstate {escaped_name}')
             if result.exit_status != 0:
                 _LOGGER.error("Failed to get VM status for '%s': %s", vm_name, result.stderr)
                 return VMState.CRASHED.value
@@ -256,7 +272,9 @@ async def start_vm(self, vm_name: str) -> bool:
                 _LOGGER.info("VM '%s' is already running", vm_name)
                 return True
 
-            result = await self.execute_command(f'virsh start "{vm_name}"')
+            # Use shlex.quote to properly escape VM names with special characters
+            escaped_name = shlex.quote(vm_name)
+            result = await self.execute_command(f'virsh start {escaped_name}')
             success = result.exit_status == 0
 
             if not success:
@@ -289,7 +307,9 @@ async def stop_vm(self, vm_name: str) -> bool:
                 _LOGGER.info("VM '%s' is already shut off", vm_name)
                 return True
 
-            result = await self.execute_command(f'virsh shutdown "{vm_name}"')
+            # Use shlex.quote to properly escape VM names with special characters
+            escaped_name = shlex.quote(vm_name)
+            result = await self.execute_command(f'virsh shutdown {escaped_name}')
             success = result.exit_status == 0
 
             if not success:
@@ -326,7 +346,8 @@ async def pause_vm(self, vm_name: str) -> bool:
                 _LOGGER.error("Cannot pause VM '%s' because it is not running (current state: %s)", vm_name, current_state)
                 return False
 
-            result = await self.execute_command(f'virsh suspend "{vm_name}"')
+            escaped_name = shlex.quote(vm_name)
+            result = await self.execute_command(f'virsh suspend {escaped_name}')
             success = result.exit_status == 0
 
             if not success:
@@ -363,7 +384,8 @@ async def resume_vm(self, vm_name: str) -> bool:
                 _LOGGER.error("Cannot resume VM '%s' because it is not paused (current state: %s)", vm_name, current_state)
                 return False
 
-            result = await self.execute_command(f'virsh resume "{vm_name}"')
+            escaped_name = shlex.quote(vm_name)
+            result = await self.execute_command(f'virsh resume {escaped_name}')
             success = result.exit_status == 0
 
             if not success:
@@ -396,7 +418,8 @@ async def restart_vm(self, vm_name: str) -> bool:
                 _LOGGER.error("Cannot restart VM '%s' because it is not running (current state: %s)", vm_name, current_state)
                 return False
 
-            result = await self.execute_command(f'virsh reboot "{vm_name}"')
+            escaped_name = shlex.quote(vm_name)
+            result = await self.execute_command(f'virsh reboot {escaped_name}')
             success = result.exit_status == 0
 
             if not success:
@@ -433,7 +456,8 @@ async def hibernate_vm(self, vm_name: str) -> bool:
                 _LOGGER.error("Cannot hibernate VM '%s' because it is not running (current state: %s)", vm_name, current_state)
                 return False
 
-            result = await self.execute_command(f'virsh dompmsuspend "{vm_name}" disk')
+            escaped_name = shlex.quote(vm_name)
+            result = await self.execute_command(f'virsh dompmsuspend {escaped_name} disk')
             success = result.exit_status == 0
 
             if not success:
@@ -466,7 +490,8 @@ async def force_stop_vm(self, vm_name: str) -> bool:
                 _LOGGER.info("VM '%s' is already shut off", vm_name)
                 return True
 
-            result = await self.execute_command(f'virsh destroy "{vm_name}"')
+            escaped_name = shlex.quote(vm_name)
+            result = await self.execute_command(f'virsh destroy {escaped_name}')
             success = result.exit_status == 0
 
             if not success:

custom_components/unraid/switch.py

@@ -150,20 +150,94 @@ def __init__(
         self._vm_name = vm_name
         self._last_known_state = None
 
-        # Remove any leading numbers and spaces for the entity ID
-        cleaned_name = ''.join(c for c in vm_name if not c.isdigit()).strip()
+        # Create a safe entity ID from VM name with collision detection
+        from .utils import normalize_name
+        safe_name = normalize_name(vm_name)
+        # Ensure it doesn't start with a number (Home Assistant requirement)
+        if safe_name and safe_name[0].isdigit():
+            safe_name = f"vm_{safe_name}"
+
+        # Handle entity ID collisions by checking existing VMs
+        safe_name = self._ensure_unique_entity_id(safe_name, vm_name, coordinator)
 
         super().__init__(
             coordinator,
             UnraidSwitchEntityDescription(
-                key=f"vm_{cleaned_name}",
+                key=f"vm_{safe_name}",
                 name=f"{vm_name}",
                 value_fn=self._get_vm_state,
             )
         )
         self._attr_entity_registry_enabled_default = True
 
         # Get OS type for specific model info
+        self._get_os_type_info(vm_name, coordinator)
+
+    def _ensure_unique_entity_id(self, base_name: str, vm_name: str, coordinator) -> str:
+        """Ensure the entity ID is unique by checking for collisions with existing VMs.
+
+        Args:
+            base_name: The normalized base name for the entity ID
+            vm_name: The original VM name
+            coordinator: The coordinator containing VM data
+
+        Returns:
+            A unique entity ID that won't collide with existing VMs
+        """
+        from .utils import normalize_name
+
+        # Get all existing VM names from coordinator data
+        existing_vms = coordinator.data.get("vms", [])
+        existing_normalized_names = set()
+
+        for vm in existing_vms:
+            existing_vm_name = vm.get("name", "")
+            if existing_vm_name and existing_vm_name != vm_name:  # Don't include current VM
+                existing_normalized = normalize_name(existing_vm_name)
+                if existing_normalized and existing_normalized[0].isdigit():
+                    existing_normalized = f"vm_{existing_normalized}"
+                existing_normalized_names.add(existing_normalized)
+
+        # If no collision, return the base name
+        if base_name not in existing_normalized_names:
+            return base_name
+
+        # Handle collision by adding a suffix based on original VM name characteristics
+        # Strategy: Use distinguishing characteristics from the original name
+
+        # Try to create a unique suffix based on the original name
+        import re
+
+        # Extract unique characteristics from the original VM name
+        # 1. Check for different separators (dash vs space vs underscore)
+        if '-' in vm_name and ' ' not in vm_name:
+            suffix = "dash"
+        elif ' ' in vm_name and '-' not in vm_name:
+            suffix = "space"
+        elif '_' in vm_name:
+            suffix = "underscore"
+        else:
+            # 2. Use position-based numbering or character-based differentiation
+            # Extract any numbers or special patterns
+            numbers = re.findall(r'\d+', vm_name)
+            if numbers:
+                suffix = f"n{numbers[-1]}"  # Use last number found
+            else:
+                # 3. Use length or hash-based suffix as last resort
+                suffix = f"len{len(vm_name)}"
+
+        candidate_name = f"{base_name}_{suffix}"
+
+        # If still collision, fall back to incremental numbering
+        counter = 2
+        while candidate_name in existing_normalized_names:
+            candidate_name = f"{base_name}_{suffix}_{counter}"
+            counter += 1
+
+        return candidate_name
+
+    def _get_os_type_info(self, vm_name: str, coordinator) -> None:
+        """Get OS type for specific model info."""
         for vm in coordinator.data.get("vms", []):
             if vm["name"] == vm_name and "os_type" in vm:
                 self._attr_device_info["model"] = f"{vm['os_type'].capitalize()} Virtual Machine"