From 08d9f78a35b0fffb1c698c245a74b4a1867bd20e Mon Sep 17 00:00:00 2001 From: Mathieu Broillet Date: Sat, 19 Oct 2024 11:26:43 +0200 Subject: [PATCH] improved ssh connections and reliability --- .../computer/__init__.py | 34 ++++++++++++++----- .../computer/ssh_client.py | 15 +++++--- .../easy_computer_manager/computer/utils.py | 3 +- .../easy_computer_manager/switch.py | 8 ++--- requirements.txt | 5 +-- 5 files changed, 45 insertions(+), 20 deletions(-) diff --git a/custom_components/easy_computer_manager/computer/__init__.py b/custom_components/easy_computer_manager/computer/__init__.py index 2377b27..f91d18e 100644 --- a/custom_components/easy_computer_manager/computer/__init__.py +++ b/custom_components/easy_computer_manager/computer/__init__.py @@ -34,8 +34,9 @@ class Computer: self._connection: SSHClient = SSHClient(host, username, password, port) asyncio.create_task(self._connection.connect(computer=self)) + # asyncio.create_task(self.update()) - async def update(self) -> None: + async def update(self, state: Optional[bool] = True, timeout: Optional[int] = 2) -> None: """Update computer details.""" async def update_operating_system(): @@ -77,11 +78,30 @@ class Computer: # TODO: implement for Windows pass - # Reconnect if connection is lost and init is already done - if self.initialized and not self._connection.is_connection_alive(): - await self._connection.connect() + if not state or not await self.is_on(): + LOGGER.debug("ECM Computer is off, skipping update") + return + else: + # Wait for connection to be established before updating or give up after timeout + for i in range(timeout * 4): + if not self._connection.is_connection_alive(): + await asyncio.sleep(0.25) + else: + break + else: + # Reconnect if connection is lost and init is already done + if self.initialized: + await self._connection.connect() + + if not self._connection.is_connection_alive(): + LOGGER.debug( + f"Computer seems ON but the SSH connection is dead (timeout={timeout}s), skipping update") + return + else: + LOGGER.debug( + f"Computer seems ON but the SSH connection is dead (timeout={timeout}s), skipping update") + return - if self._connection.is_connection_alive(): await update_operating_system() await update_operating_system_version() await update_desktop_environment() @@ -213,6 +233,4 @@ class Computer: async def run_manually(self, command: str) -> CommandOutput: """Run a custom command manually via SSH.""" - result = await self._connection.execute_command(command) - - return CommandOutput(command, result[0], result[1], result[2]) + return await self._connection.execute_command(command) diff --git a/custom_components/easy_computer_manager/computer/ssh_client.py b/custom_components/easy_computer_manager/computer/ssh_client.py index 62098c8..9b15a2c 100644 --- a/custom_components/easy_computer_manager/computer/ssh_client.py +++ b/custom_components/easy_computer_manager/computer/ssh_client.py @@ -1,8 +1,10 @@ import asyncio -import paramiko from typing import Optional +import paramiko + from custom_components.easy_computer_manager import LOGGER +from custom_components.easy_computer_manager.computer import CommandOutput class SSHClient: @@ -56,15 +58,17 @@ class SSHClient: port=self.port ) - async def execute_command(self, command: str) -> tuple[int, str, str]: + async def execute_command(self, command: str) -> CommandOutput: """Execute a command on the SSH server asynchronously.""" try: stdin, stdout, stderr = self._connection.exec_command(command) exit_status = stdout.channel.recv_exit_status() - return exit_status, stdout.read().decode(), stderr.read().decode() + return CommandOutput(command, exit_status, stdout.read().decode(), stderr.read().decode()) + except (paramiko.SSHException, EOFError) as exc: LOGGER.debug(f"Failed to execute command on {self.host}: {exc}") + return CommandOutput(command, -1, "", "") def is_connection_alive(self) -> bool: """Check if the connection is still alive asynchronously.""" @@ -75,6 +79,9 @@ class SSHClient: try: transport = self._connection.get_transport() transport.send_ignore() + + self._connection.exec_command('ls', timeout=1) return True - except EOFError: + + except Exception: return False diff --git a/custom_components/easy_computer_manager/computer/utils.py b/custom_components/easy_computer_manager/computer/utils.py index f8d22b0..89dd779 100644 --- a/custom_components/easy_computer_manager/computer/utils.py +++ b/custom_components/easy_computer_manager/computer/utils.py @@ -13,7 +13,8 @@ async def format_debug_information(computer: 'Computer'): # importing Computer 'username': computer.username, 'port': computer.port, 'dualboot': computer.dualboot, - 'is_on': await computer.is_on() + 'is_on': await computer.is_on(), + 'is_connected': computer._connection.is_connection_alive() }, 'grub': { 'windows_entry': computer.windows_entry_grub diff --git a/custom_components/easy_computer_manager/switch.py b/custom_components/easy_computer_manager/switch.py index ee54c50..096c933 100644 --- a/custom_components/easy_computer_manager/switch.py +++ b/custom_components/easy_computer_manager/switch.py @@ -125,9 +125,6 @@ class ComputerSwitch(SwitchEntity): @property def device_info(self) -> DeviceInfo | None: """Return the device information.""" - # TODO: remove this? - # if self._host is None: - # return None return DeviceInfo( identifiers={(DOMAIN, self.computer.mac)}, @@ -164,7 +161,7 @@ class ComputerSwitch(SwitchEntity): """Ping the computer to see if it is online and update the state.""" self._state = await self.computer.is_on() - await self.computer.update() + await self.computer.update(self._state) # Update the state attributes and the connection only if the computer is on if self._state: @@ -191,6 +188,7 @@ class ComputerSwitch(SwitchEntity): async def start_computer_to_windows(self) -> None: """(Service Handler) Start the computer to Windows""" + async def wait_task(): while not self.is_on: await asyncio.sleep(3) @@ -215,7 +213,7 @@ class ComputerSwitch(SwitchEntity): async def change_audio_config(self, volume: int | None = None, mute: bool | None = None, input_device: str | None = None, - output_device: str | None = None) -> None: + output_device: str | None = None) -> None: """(Service Handler) Change the audio configuration using a YAML config file.""" await self.computer.set_audio_config(volume, mute, input_device, output_device) diff --git a/requirements.txt b/requirements.txt index f8242e8..c0e2ae4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ wakeonlan~=3.1.0 -homeassistant~=2024.3.1 -asyncssh~=2.16.0 \ No newline at end of file +homeassistant +paramiko~=3.5.0 +voluptuous~=0.15.2 \ No newline at end of file