Skip to content

Commit

Permalink
Improved VM Status reporting back to HA. Added extra attributes for V…
Browse files Browse the repository at this point in the history
…Ms and Containers.
  • Loading branch information
MrD3y5eL committed Oct 22, 2024
1 parent a18eac7 commit c02676b
Show file tree
Hide file tree
Showing 5 changed files with 211 additions and 36 deletions.
1 change: 1 addition & 0 deletions custom_components/unraid/const.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""Constants for the Unraid integration."""
from homeassistant.const import Platform


DOMAIN = "unraid"
DEFAULT_PORT = 22
DEFAULT_CHECK_INTERVAL = 300 # seconds
Expand Down
10 changes: 8 additions & 2 deletions custom_components/unraid/coordinator.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,16 +32,22 @@ def __init__(self, hass: HomeAssistant, api: UnraidAPI, entry: ConfigEntry) -> N
async def _async_update_data(self) -> Dict[str, Any]:
"""Fetch data from Unraid."""
try:
# Fetch VM data first for faster switch response
vms = await self.api.get_vms()

# Then fetch the rest of the data
data = {
"vms": vms,
"system_stats": await self.api.get_system_stats(),
"docker_containers": await self.api.get_docker_containers(),
"vms": await self.api.get_vms(),
"user_scripts": await self.api.get_user_scripts(),
}

if self.has_ups:
ups_info = await self.api.get_ups_info()
if ups_info: # Only add UPS info if it's not empty
if ups_info:
data["ups_info"] = ups_info

return data
except Exception as err:
_LOGGER.error("Error communicating with Unraid: %s", err)
Expand Down
2 changes: 1 addition & 1 deletion custom_components/unraid/manifest.json
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,6 @@
"issue_tracker": "https://github.com/domalab/ha-unraid/issues",
"requirements": [],
"ssdp": [],
"version": "0.1.3",
"version": "0.1.4",
"zeroconf": []
}
58 changes: 54 additions & 4 deletions custom_components/unraid/switch.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
"""Switch platform for Unraid."""
from __future__ import annotations

from typing import Any, Dict

from homeassistant.components.switch import SwitchEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from homeassistant.exceptions import HomeAssistantError

from .const import DOMAIN
from .coordinator import UnraidDataUpdateCoordinator
Expand Down Expand Up @@ -79,9 +82,21 @@ def is_on(self) -> bool:
"""Return true if the container is running."""
for container in self.coordinator.data["docker_containers"]:
if container["name"] == self._container_name:
return container["status"].lower() == "running"
return container["state"] == "running"
return False

@property
def extra_state_attributes(self) -> Dict[str, Any]:
"""Return the state attributes."""
for container in self.coordinator.data["docker_containers"]:
if container["name"] == self._container_name:
return {
"container_id": container["id"],
"status": container["status"],
"image": container["image"]
}
return {}

async def async_turn_on(self, **kwargs) -> None:
"""Turn the container on."""
await self.coordinator.api.start_container(self._container_name)
Expand All @@ -100,7 +115,27 @@ def __init__(self, coordinator: UnraidDataUpdateCoordinator, vm_name: str) -> No
super().__init__(coordinator, f"vm_{vm_name}")
self._vm_name = vm_name
self._attr_name = f"Unraid VM {vm_name}"
self._attr_icon = "mdi:desktop-classic"
self._attr_entity_registry_enabled_default = True
self._attr_assumed_state = False

@property
def icon(self) -> str:
"""Return the icon to use for the VM."""
for vm in self.coordinator.data["vms"]:
if vm["name"] == self._vm_name:
if vm.get("os_type") == "windows":
return "mdi:microsoft-windows"
elif vm.get("os_type") == "linux":
return "mdi:linux"
return "mdi:desktop-tower"
return "mdi:desktop-tower"

@property
def available(self) -> bool:
"""Return if the switch is available."""
if not self.coordinator.last_update_success:
return False
return any(vm["name"] == self._vm_name for vm in self.coordinator.data["vms"])

@property
def is_on(self) -> bool:
Expand All @@ -109,13 +144,28 @@ def is_on(self) -> bool:
if vm["name"] == self._vm_name:
return vm["status"].lower() == "running"
return False

@property
def extra_state_attributes(self) -> Dict[str, Any]:
"""Return the state attributes."""
for vm in self.coordinator.data["vms"]:
if vm["name"] == self._vm_name:
return {
"os_type": vm.get("os_type", "unknown"),
"status": vm.get("status", "unknown"),
}
return {}

async def async_turn_on(self, **kwargs) -> None:
"""Turn the VM on."""
await self.coordinator.api.start_vm(self._vm_name)
success = await self.coordinator.api.start_vm(self._vm_name)
if not success:
raise HomeAssistantError(f"Failed to start VM {self._vm_name}")
await self.coordinator.async_request_refresh()

async def async_turn_off(self, **kwargs) -> None:
"""Turn the VM off."""
await self.coordinator.api.stop_vm(self._vm_name)
success = await self.coordinator.api.stop_vm(self._vm_name)
if not success:
raise HomeAssistantError(f"Failed to stop VM {self._vm_name}")
await self.coordinator.async_request_refresh()
176 changes: 147 additions & 29 deletions custom_components/unraid/unraid.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,53 @@
from typing import Dict, List, Any, Optional
import re
from async_timeout import timeout
from enum import Enum
import json

_LOGGER = logging.getLogger(__name__)

class VMState(Enum):
"""VM states matching Unraid/libvirt states."""
RUNNING = 'running'
STOPPED = 'shut off'
PAUSED = 'paused'
IDLE = 'idle'
IN_SHUTDOWN = 'in shutdown'
CRASHED = 'crashed'
SUSPENDED = 'pmsuspended'

@classmethod
def is_running(cls, state: str) -> bool:
"""Check if the state represents a running VM."""
return state.lower() == cls.RUNNING.value

@classmethod
def parse(cls, state: str) -> str:
"""Parse the VM state string."""
state = state.lower().strip()
try:
return next(s.value for s in cls if s.value == state)
except StopIteration:
return state

class ContainerStates(Enum):
"""Docker container states."""
RUNNING = 'running'
EXITED = 'exited'
PAUSED = 'paused'
RESTARTING = 'restarting'
DEAD = 'dead'
CREATED = 'created'

@classmethod
def parse(cls, state: str) -> str:
"""Parse the container state string."""
state = state.lower().strip()
try:
return next(s.value for s in cls if s.value == state)
except StopIteration:
return state

class UnraidAPI:
"""API client for interacting with Unraid servers."""

Expand Down Expand Up @@ -397,18 +441,35 @@ async def get_docker_containers(self) -> List[Dict[str, Any]]:
"""Fetch information about Docker containers."""
try:
_LOGGER.debug("Fetching Docker container information")
result = await self.execute_command("docker ps -a --format '{{.Names}}|{{.State}}'")
# Get basic container info with proven format
result = await self.execute_command("docker ps -a --format '{{.Names}}|{{.State}}|{{.ID}}|{{.Image}}'")
if result.exit_status != 0:
_LOGGER.error("Docker container list command failed with exit status %d", result.exit_status)
return []

containers = []
for line in result.stdout.splitlines():
parts = line.split('|')
if len(parts) == 2:
containers.append({"name": parts[0], "status": parts[1]})
if len(parts) == 4: # Now expecting 4 parts
container_name = parts[0].strip()
# Get container icon if available
icon_path = f"/var/lib/docker/unraid/images/{container_name}-icon.png"
icon_result = await self.execute_command(
f"[ -f {icon_path} ] && (base64 {icon_path}) || echo ''"
)
icon_data = icon_result.stdout[0] if icon_result.exit_status == 0 else ""

containers.append({
"name": container_name,
"state": ContainerStates.parse(parts[1].strip()),
"status": parts[1].strip(),
"id": parts[2].strip(),
"image": parts[3].strip(),
"icon": icon_data
})
else:
_LOGGER.warning("Unexpected format in docker container output: %s", line)

return containers
except Exception as e:
_LOGGER.error("Error getting docker containers: %s", str(e))
Expand All @@ -417,8 +478,8 @@ async def get_docker_containers(self) -> List[Dict[str, Any]]:
async def start_container(self, container_name: str) -> bool:
"""Start a Docker container."""
try:
_LOGGER.debug("Starting Docker container: %s", container_name)
result = await self.execute_command(f"docker start {container_name}")
_LOGGER.debug("Starting container: %s", container_name)
result = await self.execute_command(f'docker start "{container_name}"')
if result.exit_status != 0:
_LOGGER.error("Failed to start container %s: %s", container_name, result.stderr)
return False
Expand All @@ -427,12 +488,12 @@ async def start_container(self, container_name: str) -> bool:
except Exception as e:
_LOGGER.error("Error starting container %s: %s", container_name, str(e))
return False

async def stop_container(self, container_name: str) -> bool:
"""Stop a Docker container."""
try:
_LOGGER.debug("Stopping Docker container: %s", container_name)
result = await self.execute_command(f"docker stop {container_name}")
_LOGGER.debug("Stopping container: %s", container_name)
result = await self.execute_command(f'docker stop "{container_name}"')
if result.exit_status != 0:
_LOGGER.error("Failed to stop container %s: %s", container_name, result.stderr)
return False
Expand All @@ -455,43 +516,100 @@ async def get_vms(self) -> List[Dict[str, Any]]:
for line in result.stdout.splitlines():
if line.strip():
name = line.strip()
status = await self._get_vm_status(name)
vms.append({"name": name, "status": status})
status = await self.get_vm_status(name)
os_type = await self.get_vm_os_info(name)
vms.append({
"name": name,
"status": status,
"os_type": os_type
})
return vms
except Exception as e:
_LOGGER.error("Error getting VMs: %s", str(e))
return []

async def get_vm_os_info(self, vm_name: str) -> str:
"""Get the OS type of a VM."""
try:
# First try to get OS info from VM XML
result = await self.execute_command(f'virsh dumpxml "{vm_name}" | grep "<os>"')
xml_output = result.stdout

# Check for Windows-specific indicators
if any(indicator in '\n'.join(xml_output).lower() for indicator in ['windows', 'win', 'microsoft']):
return 'windows'

# Try to get detailed OS info if available
result = await self.execute_command(f'virsh domosinfo "{vm_name}" 2>/dev/null')
if result.exit_status == 0:
os_info = '\n'.join(result.stdout).lower()
if any(indicator in os_info for indicator in ['windows', 'win', 'microsoft']):
return 'windows'
elif any(indicator in os_info for indicator in ['linux', 'unix', 'ubuntu', 'debian', 'centos', 'fedora', 'rhel']):
return 'linux'

# Default to checking common paths in VM name
vm_name_lower = vm_name.lower()
if any(win_term in vm_name_lower for win_term in ['windows', 'win']):
return 'windows'
elif any(linux_term in vm_name_lower for linux_term in ['linux', 'ubuntu', 'debian', 'centos', 'fedora', 'rhel']):
return 'linux'

return 'unknown'
except Exception as e:
_LOGGER.debug("Error getting OS info for VM %s: %s", vm_name, str(e))
return 'unknown'

async def _get_vm_status(self, vm_name: str) -> str:
"""Get the status of a specific virtual machine."""
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}")
if result.exit_status != 0:
_LOGGER.error("VM status command for %s failed with exit status %d", vm_name, result.exit_status)
return "unknown"
return result.stdout.strip()
_LOGGER.error("Failed to get VM status for %s: %s", vm_name, result.stderr)
return VMState.CRASHED.value
return VMState.parse(result.stdout.strip())
except Exception as e:
_LOGGER.error("Error getting VM status for %s: %s", vm_name, str(e))
return "unknown"
return VMState.CRASHED.value

async def start_vm(self, vm_name: str) -> bool:
"""Start a virtual machine."""
async def stop_vm(self, vm_name: str) -> bool:
"""Stop a virtual machine using ACPI shutdown."""
try:
_LOGGER.debug("Starting VM: %s", vm_name)
result = await self.execute_command(f"virsh start {vm_name}")
return result.exit_status == 0 and "started" in result.stdout.lower()
_LOGGER.debug("Stopping VM: %s", vm_name)
result = await self.execute_command(f'virsh shutdown "{vm_name}" --mode acpi')
success = result.exit_status == 0

if success:
# Wait for the VM to actually shut down
for _ in range(30): # Wait up to 60 seconds
await asyncio.sleep(2)
status = await self.get_vm_status(vm_name)
if status == VMState.STOPPED.value:
return True
return False
return False
except Exception as e:
_LOGGER.error("Error starting VM %s: %s", vm_name, str(e))
_LOGGER.error("Error stopping VM %s: %s", vm_name, str(e))
return False

async def stop_vm(self, vm_name: str) -> bool:
"""Stop a virtual machine."""
async def start_vm(self, vm_name: str) -> bool:
"""Start a virtual machine and wait for it to be running."""
try:
_LOGGER.debug("Stopping VM: %s", vm_name)
result = await self.execute_command(f"virsh shutdown {vm_name}")
return result.exit_status == 0 and "shutting down" in result.stdout.lower()
_LOGGER.debug("Starting VM: %s", vm_name)
result = await self.execute_command(f'virsh start "{vm_name}"')
success = result.exit_status == 0

if success:
# Wait for the VM to actually start
for _ in range(15): # Wait up to 30 seconds
await asyncio.sleep(2)
status = await self.get_vm_status(vm_name)
if status == VMState.RUNNING.value:
return True
return False
return False
except Exception as e:
_LOGGER.error("Error stopping VM %s: %s", vm_name, str(e))
_LOGGER.error("Error starting VM %s: %s", vm_name, str(e))
return False

async def get_user_scripts(self) -> List[Dict[str, Any]]:
Expand Down

0 comments on commit c02676b

Please sign in to comment.