From bc6b7e4ae993fe2a3a282de94798069292049f8c Mon Sep 17 00:00:00 2001 From: Qubasa Date: Mon, 23 Jun 2025 14:21:51 +0200 Subject: [PATCH] clan-lib: Rename check_machine_online to can_ssh_login. Move to Remote object --- pkgs/clan-cli/clan_lib/api/network.py | 41 ------------- pkgs/clan-cli/clan_lib/ssh/remote.py | 65 ++++++++++++++++++--- pkgs/clan-cli/clan_lib/tests/test_create.py | 5 +- 3 files changed, 58 insertions(+), 53 deletions(-) delete mode 100644 pkgs/clan-cli/clan_lib/api/network.py diff --git a/pkgs/clan-cli/clan_lib/api/network.py b/pkgs/clan-cli/clan_lib/api/network.py deleted file mode 100644 index 13c21d8be..000000000 --- a/pkgs/clan-cli/clan_lib/api/network.py +++ /dev/null @@ -1,41 +0,0 @@ -import logging -import time -from dataclasses import dataclass -from typing import Literal - -from clan_lib.api import API -from clan_lib.cmd import RunOpts -from clan_lib.errors import ClanError -from clan_lib.ssh.remote import Remote - -log = logging.getLogger(__name__) - - -@dataclass -class ConnectionOptions: - timeout: int = 2 - retries: int = 10 - - -@API.register -def check_machine_online( - remote: Remote, opts: ConnectionOptions | None = None -) -> Literal["Online", "Offline"]: - timeout = opts.timeout if opts and opts.timeout else 2 - - for _ in range(opts.retries if opts and opts.retries else 10): - with remote.ssh_control_master() as ssh: - res = ssh.run( - ["true"], - RunOpts(timeout=timeout, check=False, needs_user_terminal=True), - ) - - if res.returncode == 0: - return "Online" - - if "Host key verification failed." in res.stderr: - raise ClanError(res.stderr.strip()) - - time.sleep(timeout) - - return "Offline" diff --git a/pkgs/clan-cli/clan_lib/ssh/remote.py b/pkgs/clan-cli/clan_lib/ssh/remote.py index 8c73dbdc8..a547e7367 100644 --- a/pkgs/clan-cli/clan_lib/ssh/remote.py +++ b/pkgs/clan-cli/clan_lib/ssh/remote.py @@ -6,16 +6,19 @@ import shlex import socket import subprocess import sys +import time from collections.abc import Iterator from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path from shlex import quote from tempfile import TemporaryDirectory +from typing import Literal from clan_cli.ssh.host_key import HostKeyCheck -from clan_lib.cmd import CmdOut, RunOpts, run +from clan_lib.api import API +from clan_lib.cmd import ClanCmdError, ClanCmdTimeoutError, CmdOut, RunOpts, run from clan_lib.colors import AnsiColor from clan_lib.errors import ClanError # Assuming these are available from clan_lib.nix import nix_shell @@ -434,12 +437,56 @@ class Remote: self.check_sshpass_errorcode(res) def is_ssh_reachable(self) -> bool: - address_family = socket.AF_INET6 if ":" in self.address else socket.AF_INET - with socket.socket(address_family, socket.SOCK_STREAM) as sock: - sock.settimeout(2) - try: - sock.connect((self.address, self.port or 22)) - except OSError: - return False + return is_ssh_reachable(self) - return True + +@dataclass(frozen=True) +class ConnectionOptions: + timeout: int = 2 + retries: int = 10 + + +@API.register +def can_ssh_login( + remote: Remote, opts: ConnectionOptions | None = None +) -> Literal["Online", "Offline"]: + if opts is None: + opts = ConnectionOptions() + + for _ in range(opts.retries): + with remote.ssh_control_master() as ssh: + try: + res = ssh.run( + ["true"], + RunOpts(timeout=opts.timeout, needs_user_terminal=True), + ) + return "Online" + except ClanCmdTimeoutError: + pass + except ClanCmdError as e: + if "Host key verification failed." in e.cmd.stderr: + raise ClanError(res.stderr.strip()) from e + else: + time.sleep(opts.timeout) + + return "Offline" + + +@API.register +def is_ssh_reachable(remote: Remote, opts: ConnectionOptions | None = None) -> bool: + if opts is None: + opts = ConnectionOptions() + + address_family = socket.AF_INET6 if ":" in remote.address else socket.AF_INET + for _ in range(opts.retries): + with socket.socket(address_family, socket.SOCK_STREAM) as sock: + sock.settimeout(opts.timeout) + try: + sock.connect((remote.address, remote.port or 22)) + return True + except (TimeoutError, OSError): + pass + else: + time.sleep(opts.timeout) + + return False diff --git a/pkgs/clan-cli/clan_lib/tests/test_create.py b/pkgs/clan-cli/clan_lib/tests/test_create.py index 4c8b7fd0c..a3a3a185d 100644 --- a/pkgs/clan-cli/clan_lib/tests/test_create.py +++ b/pkgs/clan-cli/clan_lib/tests/test_create.py @@ -18,7 +18,6 @@ from clan_cli.vars.generate import generate_vars_for_machine, get_generators_clo from clan_lib.api.disk import hw_main_disk_options, set_machine_disk_schema from clan_lib.api.modules import list_modules -from clan_lib.api.network import check_machine_online from clan_lib.cmd import RunOpts, run from clan_lib.dirs import specific_machine_dir from clan_lib.errors import ClanError @@ -34,7 +33,7 @@ from clan_lib.nix_models.clan import ( ) from clan_lib.nix_models.clan import InventoryMachineDeploy as MachineDeploy from clan_lib.persist.util import set_value_by_path -from clan_lib.ssh.remote import Remote +from clan_lib.ssh.remote import Remote, can_ssh_login log = logging.getLogger(__name__) @@ -201,7 +200,7 @@ def test_clan_create_api( target_host = machine.target_host().override( private_key=private_key, host_key_check=HostKeyCheck.NONE ) - result = check_machine_online(target_host) + result = can_ssh_login(target_host) assert result == "Online", f"Machine {machine.name} is not online" ssh_keys = [