improved ssh connections and reliability

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)

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

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

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)
 

requirements.txt

@@ -1,3 +1,4 @@
 wakeonlan~=3.1.0
-homeassistant~=2024.3.1
-asyncssh~=2.16.0
+homeassistant
+paramiko~=3.5.0
+voluptuous~=0.15.2