HA-EasyComputerManager/custom_components/easy_computer_manage/switch.py

# Some code is from the official wake_on_lan integration

from __future__ import annotations

import asyncio
import logging
import subprocess as sp
from typing import Any

import voluptuous as vol
import wakeonlan
from paramiko.ssh_exception import NoValidConnectionsError, SSHException, AuthenticationException

from homeassistant.components.switch import (
    PLATFORM_SCHEMA as PARENT_PLATFORM_SCHEMA,
    SwitchEntity,
)
from homeassistant.const import (
    CONF_BROADCAST_ADDRESS,
    CONF_BROADCAST_PORT,
    CONF_HOST,
    CONF_MAC,
    CONF_NAME,
    CONF_PASSWORD,
    CONF_PORT,
    CONF_USERNAME,
)
from homeassistant.core import HomeAssistant
from homeassistant.helpers import (
    config_validation as cv,
    device_registry as dr,
    entity_platform,
)
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import utils
from .const import SERVICE_RESTART_TO_WINDOWS_FROM_LINUX, SERVICE_PUT_COMPUTER_TO_SLEEP, \
    SERVICE_START_COMPUTER_TO_WINDOWS
from homeassistant.config_entries import ConfigEntry

_LOGGER = logging.getLogger(__name__)

CONF_OFF_ACTION = "turn_off"

DEFAULT_NAME = "Computer Management (WoL, SoL)"
DEFAULT_PING_TIMEOUT = 1

PLATFORM_SCHEMA = PARENT_PLATFORM_SCHEMA.extend(
    {
        vol.Required(CONF_MAC): cv.string,
        vol.Optional(CONF_BROADCAST_ADDRESS): cv.string,
        vol.Optional(CONF_BROADCAST_PORT): cv.port,
        vol.Required(CONF_HOST): cv.string,
        vol.Required("dualboot", default=False): cv.boolean,
        vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
        vol.Required(CONF_USERNAME, default="root"): cv.string,
        vol.Required(CONF_PASSWORD, default="root"): cv.string,
        vol.Optional(CONF_PORT, default=22): cv.string,
    }
)


async def async_setup_entry(
        hass: HomeAssistant,
        config: ConfigEntry,
        async_add_entities: AddEntitiesCallback,
) -> None:
    mac_address: str = config.data.get(CONF_MAC)
    broadcast_address: str | None = config.data.get(CONF_BROADCAST_ADDRESS)
    broadcast_port: int | None = config.data.get(CONF_BROADCAST_PORT)
    host: str = config.data.get(CONF_HOST)
    name: str = config.data.get(CONF_NAME)
    dualboot: bool = config.data.get("dualboot")
    username: str = config.data.get(CONF_USERNAME)
    password: str = config.data.get(CONF_PASSWORD)
    port: int | None = config.data.get(CONF_PORT)

    async_add_entities(
        [
            ComputerSwitch(
                hass,
                name,
                host,
                mac_address,
                broadcast_address,
                broadcast_port,
                dualboot,
                username,
                password,
                port,
            ),
        ],
        host is not None,
    )

    platform = entity_platform.async_get_current_platform()
    platform.async_register_entity_service(
        SERVICE_RESTART_TO_WINDOWS_FROM_LINUX,
        {},
        SERVICE_RESTART_TO_WINDOWS_FROM_LINUX,
    )
    platform.async_register_entity_service(
        SERVICE_PUT_COMPUTER_TO_SLEEP,
        {},
        SERVICE_PUT_COMPUTER_TO_SLEEP,
    )
    platform.async_register_entity_service(
        SERVICE_START_COMPUTER_TO_WINDOWS,
        {},
        SERVICE_START_COMPUTER_TO_WINDOWS,
    )


class ComputerSwitch(SwitchEntity):
    """Representation of a computer switch."""

    def __init__(
            self,
            hass: HomeAssistant,
            name: str,
            host: str | None,
            mac_address: str,
            broadcast_address: str | None,
            broadcast_port: int | None,
            dualboot: bool | False,
            username: str,
            password: str,
            port: int | None,
    ) -> None:
        """Initialize the WOL switch."""

        self._hass = hass
        self._attr_name = name
        self._host = host
        self._mac_address = mac_address
        self._broadcast_address = broadcast_address
        self._broadcast_port = broadcast_port
        self._dualboot = dualboot
        self._username = username
        self._password = password
        self._port = port
        self._state = False
        self._attr_assumed_state = host is None
        self._attr_should_poll = bool(not self._attr_assumed_state)
        self._attr_unique_id = dr.format_mac(mac_address)
        self._attr_extra_state_attributes = {}
        self._connection = utils.create_ssh_connection(self._host, self._username, self._password)

    @property
    def is_on(self) -> bool:
        """Return true if the computer switch is on."""
        return self._state

    def turn_on(self, **kwargs: Any) -> None:
        """Turn the computer on using wake on lan."""
        service_kwargs: dict[str, Any] = {}
        if self._broadcast_address is not None:
            service_kwargs["ip_address"] = self._broadcast_address
        if self._broadcast_port is not None:
            service_kwargs["port"] = self._broadcast_port

        _LOGGER.info(
            "Send magic packet to mac %s (broadcast: %s, port: %s)",
            self._mac_address,
            self._broadcast_address,
            self._broadcast_port,
        )

        wakeonlan.send_magic_packet(self._mac_address, **service_kwargs)

        if self._attr_assumed_state:
            self._state = True
            self.async_write_ha_state()

    def turn_off(self, **kwargs: Any) -> None:
        """Turn the computer off using appropriate shutdown command based on running OS and/or distro."""
        utils.shutdown_system(self._connection)

        if self._attr_assumed_state:
            self._state = False
            self.async_write_ha_state()

    def restart_to_windows_from_linux(self) -> None:
        """Restart the computer to Windows from a running Linux by setting grub-reboot and restarting."""

        if self._dualboot:
            utils.restart_to_windows_from_linux(self._connection)
        else:
            _LOGGER.error("This computer is not running a dualboot system.")

    def put_computer_to_sleep(self) -> None:
        """Put the computer to sleep using appropriate sleep command based on running OS and/or distro."""
        utils.sleep_system(self._connection)

    def start_computer_to_windows(self) -> None:
        """Start the computer to Linux, wait for it to boot, and then set grub-reboot and restart."""
        self.turn_on()

        if self._dualboot:
            # Wait for the computer to boot using a dedicated thread to avoid blocking the main thread

            self._hass.loop.create_task(self.reboot_computer_to_windows_when_on())

        else:
            _LOGGER.error("This computer is not running a dualboot system.")

    async def reboot_computer_to_windows_when_on(self) -> None:
        """Method to be run in a separate thread to wait for the computer to boot and then reboot to Windows."""
        while not self.is_on:
            await asyncio.sleep(3)

        await utils.restart_to_windows_from_linux(self._connection)

    def update(self) -> None:
        """Ping the computer to see if it is online and update the state."""
        ping_cmd = [
            "ping",
            "-c",
            "1",
            "-W",
            str(DEFAULT_PING_TIMEOUT),
            str(self._host),
        ]
        status = sp.call(ping_cmd, stdout=sp.DEVNULL, stderr=sp.DEVNULL)
        self._state = not bool(status)

        # Update the state attributes and the connection only if the computer is on
        if self._state:
            if not utils.test_connection(self._connection):
                _LOGGER.info("RENEWING SSH CONNECTION")

                if self._connection is not None:
                    self._connection.close()

                self._connection = utils.create_ssh_connection(self._host, self._username, self._password)

                try:
                    self._connection.open()
                except NoValidConnectionsError and SSHException as error:
                    _LOGGER.error("Could not connect to %s: %s", self._host, error)
                    self._state = False
                    return
                except AuthenticationException as error:
                    _LOGGER.error("Could not authenticate to %s: %s", self._host, error)
                    self._state = False
                    return

            self._attr_extra_state_attributes = {
                "operating_system": utils.get_operating_system(self._connection),
                "operating_system_version": utils.get_operating_system_version(
                    self._connection
                ),
                "mac_address": self._mac_address,
                "ip_address": self._host,
            }