custom_components/easy_computer_manager/computer/__init__.py

@@ -1,6 +1,8 @@
 import asyncio
 from typing import Optional, Dict, Any
 
+import asyncssh
+from asyncssh import SSHClientConnection
 from wakeonlan import send_magic_packet
 
 from custom_components.easy_computer_manager import const, LOGGER
@@ -8,15 +10,12 @@ from custom_components.easy_computer_manager.computer.common import OSType, Comm
 from custom_components.easy_computer_manager.computer.formatter import format_gnome_monitors_args, format_pactl_commands
 from custom_components.easy_computer_manager.computer.parser import parse_gnome_monitors_output, parse_pactl_output, \
     parse_bluetoothctl
-from custom_components.easy_computer_manager.computer.ssh_client import SSHClient
 
 
 class Computer:
     def __init__(self, host: str, mac: str, username: str, password: str, port: int = 22,
                  dualboot: bool = False) -> None:
         """Initialize the Computer object."""
-        self.initialized = False  # used to avoid duplicated ssh connections
-
         self.host = host
         self.mac = mac
         self.username = username
@@ -32,12 +31,38 @@ class Computer:
         self.audio_config: Dict[str, Optional[Dict]] = {}
         self.bluetooth_devices: Dict[str, Any] = {}
 
-        self._connection: SSHClient = SSHClient(host, username, password, port)
+        self._connection: Optional[SSHClientConnection] = None
-        asyncio.create_task(self._connection.connect(computer=self))
+
+        asyncio.create_task(self.update())
+
+    async def _connect(self, retried: bool = False) -> None:
+        """Open an asynchronous SSH connection."""
+        try:
+            client = await asyncssh.connect(
+                self.host,
+                username=self.username,
+                password=self._password,
+                port=self.port,
+                known_hosts=None
+            )
+            asyncssh.set_log_level("ERROR")
+            self._connection = client
+        except (OSError, asyncssh.Error) as exc:
+            if retried:
+                await self._connect(retried=True)
+            else:
+                raise ValueError(f"Failed to connect to {self.host}: {exc}")
+
+    async def _renew_connection(self) -> None:
+        """Renew the SSH connection if it is closed."""
+        if self._connection is None or self._connection.is_closed:
+            self._connection = await self._connect()
 
     async def update(self) -> None:
         """Update computer details."""
 
+        await self._renew_connection()
+
         async def update_operating_system():
             self.operating_system = await self._detect_operating_system()
 
@@ -77,18 +102,13 @@ class Computer:
                 # TODO: implement for Windows
                 pass
 
-        # Reconnect if connection is lost and init is already done
+        await update_operating_system()
-        if self.initialized and not self._connection.is_connection_alive():
+        await update_operating_system_version()
-            await self._connection.connect()
+        await update_desktop_environment()
-
+        await update_windows_entry_grub()
-        if self._connection.is_connection_alive():
+        await update_monitors_config()
-            await update_operating_system()
+        await update_audio_config()
-            await update_operating_system_version()
+        await update_bluetooth_devices()
-            await update_desktop_environment()
-            await update_windows_entry_grub()
-            await update_monitors_config()
-            await update_audio_config()
-            await update_bluetooth_devices()
 
     async def _detect_operating_system(self) -> OSType:
         """Detect the operating system of the computer."""
@@ -213,6 +233,9 @@ class Computer:
 
     async def run_manually(self, command: str) -> CommandOutput:
         """Run a custom command manually via SSH."""
-        result = await self._connection.execute_command(command)
+        if not self._connection:
+            await self._connect()
 
-        return CommandOutput(command, result[0], result[1], result[2])
+        result = await self._connection.run(command)
+
+        return CommandOutput(command, result.exit_status, result.stdout, result.stderr)

custom_components/easy_computer_manager/computer/ssh_client.py

@@ -1,84 +0,0 @@
-import asyncio
-from typing import Optional
-
-import paramiko
-
-from custom_components.easy_computer_manager import LOGGER
-
-
-class SSHClient:
-    def __init__(self, host, username, password, port):
-        self.host = host
-        self.username = username
-        self._password = password
-        self.port = port
-        self._connection = None
-
-    async def connect(self, retried: bool = False, computer: Optional['Computer'] = None) -> None:
-        """Open an SSH connection using Paramiko asynchronously."""
-        self.disconnect()
-
-        loop = asyncio.get_running_loop()
-
-        try:
-            # Create the SSH client
-            client = paramiko.SSHClient()
-
-            # Set missing host key policy to automatically accept unknown host keys
-            # client.load_system_host_keys()
-            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-
-            # Offload the blocking connect call to a thread
-            await loop.run_in_executor(None, self._blocking_connect, client)
-            self._connection = client
-
-        except (OSError, paramiko.SSHException) as exc:
-            if retried:
-                await self.connect(retried=True)
-            else:
-                LOGGER.debug(f"Failed to connect to {self.host}: {exc}")
-        finally:
-            if computer is not None:
-                if hasattr(computer, "initialized"):
-                    computer.initialized = True
-
-    def disconnect(self) -> None:
-        """Close the SSH connection."""
-        if self._connection is not None:
-            self._connection.close()
-            self._connection = None
-
-    def _blocking_connect(self, client):
-        """Perform the blocking SSH connection using Paramiko."""
-        client.connect(
-            self.host,
-            username=self.username,
-            password=self._password,
-            port=self.port
-        )
-
-    async def execute_command(self, command: str) -> tuple[int, str, str]:
-        """Execute a command on the SSH server asynchronously."""
-        # if not self.is_connection_alive():
-        #     await self.connect()
-
-        try:
-            stdin, stdout, stderr = self._connection.exec_command(command)
-            exit_status = stdout.channel.recv_exit_status()
-        except (paramiko.SSHException, EOFError) as exc:
-            raise ValueError(f"Failed to execute command on {self.host}: {exc}")
-
-        return exit_status, stdout.read().decode(), stderr.read().decode()
-
-    def is_connection_alive(self) -> bool:
-        """Check if the connection is still alive asynchronously."""
-        # use the code below if is_active() returns True
-        try:
-            if self._connection is None:
-                return False
-
-            transport = self._connection.get_transport()
-            transport.send_ignore()
-            return True
-        except EOFError:
-            return False