From 94a65178c0604bb15015b4c9e55779eb7427d2fb Mon Sep 17 00:00:00 2001 From: Mathieu Broillet Date: Thu, 17 Oct 2024 22:11:10 +0200 Subject: [PATCH] fixed buggy ssh connections --- custom_components/__init__.py | 0 .../computer/__init__.py | 61 +++++--------- .../computer/ssh_client.py | 84 +++++++++++++++++++ 3 files changed, 103 insertions(+), 42 deletions(-) delete mode 100644 custom_components/__init__.py create mode 100644 custom_components/easy_computer_manager/computer/ssh_client.py diff --git a/custom_components/__init__.py b/custom_components/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/easy_computer_manager/computer/__init__.py b/custom_components/easy_computer_manager/computer/__init__.py index 1913a5a..2377b27 100644 --- a/custom_components/easy_computer_manager/computer/__init__.py +++ b/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,13 +77,18 @@ class Computer: # TODO: implement for Windows pass - await update_operating_system() - await update_operating_system_version() - await update_desktop_environment() - await update_windows_entry_grub() - await update_monitors_config() - await update_audio_config() - await update_bluetooth_devices() + # 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() + 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.""" @@ -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]) diff --git a/custom_components/easy_computer_manager/computer/ssh_client.py b/custom_components/easy_computer_manager/computer/ssh_client.py new file mode 100644 index 0000000..ab364ed --- /dev/null +++ b/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 -- 2.39.5