custom_components/easy_computer_manager/computer/__init__.py

@@ -1,8 +1,6 @@
 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
@@ -10,12 +8,15 @@ 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
@@ -31,38 +32,12 @@ class Computer:
         self.audio_config: Dict[str, Optional[Dict]] = {}
         self.bluetooth_devices: Dict[str, Any] = {}
 
-        self._connection: Optional[SSHClientConnection] = None
-
-        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()
+        self._connection: SSHClient = SSHClient(host, username, password, port)
+        asyncio.create_task(self._connection.connect(computer=self))
 
     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()
 
@@ -102,6 +77,11 @@ 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 self._connection.is_connection_alive():
             await update_operating_system()
             await update_operating_system_version()
             await update_desktop_environment()
@@ -233,9 +213,6 @@ class Computer:
 
     async def run_manually(self, command: str) -> CommandOutput:
         """Run a custom command manually via SSH."""
-        if not self._connection:
-            await self._connect()
+        result = await self._connection.execute_command(command)
 
-        result = await self._connection.run(command)
-
-        return CommandOutput(command, result.exit_status, result.stdout, result.stderr)
+        return CommandOutput(command, result[0], result[1], result[2])

custom_components/easy_computer_manager/computer/ssh_client.py

@@ -0,0 +1,84 @@
+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