Merge pull request 'fixed buggy ssh connections' (!1) from v2 into main
Reviewed-on: #1
This commit is contained in:
commit
daf2d6d63e
@ -1,8 +1,6 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
from typing import Optional, Dict, Any
|
from typing import Optional, Dict, Any
|
||||||
|
|
||||||
import asyncssh
|
|
||||||
from asyncssh import SSHClientConnection
|
|
||||||
from wakeonlan import send_magic_packet
|
from wakeonlan import send_magic_packet
|
||||||
|
|
||||||
from custom_components.easy_computer_manager import const, LOGGER
|
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.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, \
|
from custom_components.easy_computer_manager.computer.parser import parse_gnome_monitors_output, parse_pactl_output, \
|
||||||
parse_bluetoothctl
|
parse_bluetoothctl
|
||||||
|
from custom_components.easy_computer_manager.computer.ssh_client import SSHClient
|
||||||
|
|
||||||
|
|
||||||
class Computer:
|
class Computer:
|
||||||
def __init__(self, host: str, mac: str, username: str, password: str, port: int = 22,
|
def __init__(self, host: str, mac: str, username: str, password: str, port: int = 22,
|
||||||
dualboot: bool = False) -> None:
|
dualboot: bool = False) -> None:
|
||||||
"""Initialize the Computer object."""
|
"""Initialize the Computer object."""
|
||||||
|
self.initialized = False # used to avoid duplicated ssh connections
|
||||||
|
|
||||||
self.host = host
|
self.host = host
|
||||||
self.mac = mac
|
self.mac = mac
|
||||||
self.username = username
|
self.username = username
|
||||||
@ -31,38 +32,12 @@ class Computer:
|
|||||||
self.audio_config: Dict[str, Optional[Dict]] = {}
|
self.audio_config: Dict[str, Optional[Dict]] = {}
|
||||||
self.bluetooth_devices: Dict[str, Any] = {}
|
self.bluetooth_devices: Dict[str, Any] = {}
|
||||||
|
|
||||||
self._connection: Optional[SSHClientConnection] = None
|
self._connection: SSHClient = SSHClient(host, username, password, port)
|
||||||
|
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:
|
async def update(self) -> None:
|
||||||
"""Update computer details."""
|
"""Update computer details."""
|
||||||
|
|
||||||
await self._renew_connection()
|
|
||||||
|
|
||||||
async def update_operating_system():
|
async def update_operating_system():
|
||||||
self.operating_system = await self._detect_operating_system()
|
self.operating_system = await self._detect_operating_system()
|
||||||
|
|
||||||
@ -102,13 +77,18 @@ class Computer:
|
|||||||
# TODO: implement for Windows
|
# TODO: implement for Windows
|
||||||
pass
|
pass
|
||||||
|
|
||||||
await update_operating_system()
|
# Reconnect if connection is lost and init is already done
|
||||||
await update_operating_system_version()
|
if self.initialized and not self._connection.is_connection_alive():
|
||||||
await update_desktop_environment()
|
await self._connection.connect()
|
||||||
await update_windows_entry_grub()
|
|
||||||
await update_monitors_config()
|
if self._connection.is_connection_alive():
|
||||||
await update_audio_config()
|
await update_operating_system()
|
||||||
await update_bluetooth_devices()
|
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:
|
async def _detect_operating_system(self) -> OSType:
|
||||||
"""Detect the operating system of the computer."""
|
"""Detect the operating system of the computer."""
|
||||||
@ -233,9 +213,6 @@ 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."""
|
||||||
if not self._connection:
|
result = await self._connection.execute_command(command)
|
||||||
await self._connect()
|
|
||||||
|
|
||||||
result = await self._connection.run(command)
|
return CommandOutput(command, result[0], result[1], result[2])
|
||||||
|
|
||||||
return CommandOutput(command, result.exit_status, result.stdout, result.stderr)
|
|
||||||
|
@ -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
|
Loading…
Reference in New Issue
Block a user