improved ssh connections and reliability

This commit is contained in:
Mathieu Broillet 2024-10-19 11:26:43 +02:00
parent 591396895d
commit 08d9f78a35
Signed by: mathieu
GPG Key ID: A08E484FE95074C1
5 changed files with 45 additions and 20 deletions

View File

@ -34,8 +34,9 @@ class Computer:
self._connection: SSHClient = SSHClient(host, username, password, port) self._connection: SSHClient = SSHClient(host, username, password, port)
asyncio.create_task(self._connection.connect(computer=self)) 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.""" """Update computer details."""
async def update_operating_system(): async def update_operating_system():
@ -77,11 +78,30 @@ class Computer:
# TODO: implement for Windows # TODO: implement for Windows
pass pass
# Reconnect if connection is lost and init is already done if not state or not await self.is_on():
if self.initialized and not self._connection.is_connection_alive(): LOGGER.debug("ECM Computer is off, skipping update")
await self._connection.connect() 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()
await update_operating_system_version() await update_operating_system_version()
await update_desktop_environment() await update_desktop_environment()
@ -213,6 +233,4 @@ class Computer:
async def run_manually(self, command: str) -> CommandOutput: async def run_manually(self, command: str) -> CommandOutput:
"""Run a custom command manually via SSH.""" """Run a custom command manually via SSH."""
result = await self._connection.execute_command(command) return await self._connection.execute_command(command)
return CommandOutput(command, result[0], result[1], result[2])

View File

@ -1,8 +1,10 @@
import asyncio import asyncio
import paramiko
from typing import Optional from typing import Optional
import paramiko
from custom_components.easy_computer_manager import LOGGER from custom_components.easy_computer_manager import LOGGER
from custom_components.easy_computer_manager.computer import CommandOutput
class SSHClient: class SSHClient:
@ -56,15 +58,17 @@ class SSHClient:
port=self.port 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.""" """Execute a command on the SSH server asynchronously."""
try: try:
stdin, stdout, stderr = self._connection.exec_command(command) stdin, stdout, stderr = self._connection.exec_command(command)
exit_status = stdout.channel.recv_exit_status() 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: except (paramiko.SSHException, EOFError) as exc:
LOGGER.debug(f"Failed to execute command on {self.host}: {exc}") LOGGER.debug(f"Failed to execute command on {self.host}: {exc}")
return CommandOutput(command, -1, "", "")
def is_connection_alive(self) -> bool: def is_connection_alive(self) -> bool:
"""Check if the connection is still alive asynchronously.""" """Check if the connection is still alive asynchronously."""
@ -75,6 +79,9 @@ class SSHClient:
try: try:
transport = self._connection.get_transport() transport = self._connection.get_transport()
transport.send_ignore() transport.send_ignore()
self._connection.exec_command('ls', timeout=1)
return True return True
except EOFError:
except Exception:
return False return False

View File

@ -13,7 +13,8 @@ async def format_debug_information(computer: 'Computer'): # importing Computer
'username': computer.username, 'username': computer.username,
'port': computer.port, 'port': computer.port,
'dualboot': computer.dualboot, 'dualboot': computer.dualboot,
'is_on': await computer.is_on() 'is_on': await computer.is_on(),
'is_connected': computer._connection.is_connection_alive()
}, },
'grub': { 'grub': {
'windows_entry': computer.windows_entry_grub 'windows_entry': computer.windows_entry_grub

View File

@ -125,9 +125,6 @@ class ComputerSwitch(SwitchEntity):
@property @property
def device_info(self) -> DeviceInfo | None: def device_info(self) -> DeviceInfo | None:
"""Return the device information.""" """Return the device information."""
# TODO: remove this?
# if self._host is None:
# return None
return DeviceInfo( return DeviceInfo(
identifiers={(DOMAIN, self.computer.mac)}, 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.""" """Ping the computer to see if it is online and update the state."""
self._state = await self.computer.is_on() 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 # Update the state attributes and the connection only if the computer is on
if self._state: if self._state:
@ -191,6 +188,7 @@ class ComputerSwitch(SwitchEntity):
async def start_computer_to_windows(self) -> None: async def start_computer_to_windows(self) -> None:
"""(Service Handler) Start the computer to Windows""" """(Service Handler) Start the computer to Windows"""
async def wait_task(): async def wait_task():
while not self.is_on: while not self.is_on:
await asyncio.sleep(3) 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, async def change_audio_config(self, volume: int | None = None, mute: bool | None = None,
input_device: str | 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.""" """(Service Handler) Change the audio configuration using a YAML config file."""
await self.computer.set_audio_config(volume, mute, input_device, output_device) await self.computer.set_audio_config(volume, mute, input_device, output_device)

View File

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