Add SSH host-key probing and pinning, moved from the Home Assistant integration

docs/ssh.md

@@ -25,6 +25,27 @@ async with NanoKVMSSH(
 - Failures raise `NanoKVMSSHAuthenticationError`, `NanoKVMSSHConnectionError`
   (including host-key errors) or `NanoKVMSSHCommandError`.
 
+## Trusting a new device
+
+`probe_host_key()` reads the device's host key from the SSH handshake, without
+logging in. Show its fingerprint to the user, and once they approve it, pin the
+key with `host_key=`:
+
+```python
+from nanokvm.ssh_client import NanoKVMSSH, probe_host_key
+
+key = await probe_host_key("kvm.local")
+print(key.fingerprint)  # SHA256:..., compare with the device
+
+async with NanoKVMSSH("kvm.local", host_key=key.known_hosts_line) as ssh:
+    await ssh.authenticate("password")
+```
+
+A pinned key replaces the system known-hosts file, and it applies to the host
+you connect to whatever names its line carries, so a device that changed
+address keeps its key. An invalid line raises `ValueError`; a different server
+key raises `NanoKVMSSHConnectionError`.
+
 ## Metrics and watchdog
 
 ```python

nanokvm/ssh_client.py

@@ -8,10 +8,12 @@
 from dataclasses import dataclass, field
 from datetime import UTC, datetime
 from os import PathLike
+import socket
 import threading
 from types import TracebackType
 
 import paramiko
+from paramiko.hostkeys import HostKeyEntry, InvalidHostKey
 
 from .exceptions import NanoKVMError, NanoKVMNotSupportedError
 
@@ -60,6 +62,46 @@ class SSHMetrics:
     """None when the firmware has no watchdog (the Pro, or older non-Pro firmware)."""
 
 
+@dataclass(frozen=True, slots=True)
+class SSHHostKey:
+    """A server's public host key."""
+
+    host: str
+    known_hosts_line: str
+    """The key as a known-hosts line, for ``NanoKVMSSH(host_key=...)``."""
+    fingerprint: str
+    """OpenSSH-style SHA-256 fingerprint, such as ``SHA256:...``."""
+
+
+async def probe_host_key(
+    host: str, port: int = 22, *, timeout: float | None = 10
+) -> SSHHostKey:
+    """Read a server's host key from the SSH handshake, without logging in.
+
+    This is useful for establishing an initial trust-on-first-use pin with
+    ``NanoKVMSSH(host, host_key=...)``.
+    """
+
+    def probe() -> SSHHostKey:
+        sock = socket.create_connection((host, port), timeout=timeout)
+        with sock, paramiko.Transport(sock) as transport:
+            if timeout is not None:
+                transport.banner_timeout = timeout
+            transport.start_client(timeout=timeout)
+            key = transport.get_remote_server_key()
+        name = host if port == 22 else f"[{host}]:{port}"
+        return SSHHostKey(
+            host=host,
+            known_hosts_line=f"{name} {key.get_name()} {key.get_base64()}",
+            fingerprint=key.fingerprint,
+        )
+
+    try:
+        return await asyncio.to_thread(probe)
+    except (paramiko.SSHException, OSError) as e:
+        raise NanoKVMSSHConnectionError(f"SSH host-key probe failed: {e}") from e
+
+
 @dataclass
 class _CommandExecution:
     """Channel and cancellation state shared with one executor worker."""
@@ -86,6 +128,7 @@ def __init__(
         port: int = 22,
         *,
         known_hosts: str | PathLike[str] | None = None,
+        host_key: str | None = None,
         allow_unknown_host_key: bool = False,
         connect_timeout: float | None = 10,
         banner_timeout: float | None = 10,
@@ -98,6 +141,7 @@ def __init__(
         self.port = port
         self.username = username
         self.known_hosts = known_hosts
+        self.host_key = _parse_host_key(host_key) if host_key is not None else None
         self.allow_unknown_host_key = allow_unknown_host_key
         self.connect_timeout = connect_timeout
         self.banner_timeout = banner_timeout
@@ -132,7 +176,15 @@ async def authenticate(self, password: str) -> None:
         def connect() -> None:
             # Host-key files are read here so the event loop never blocks on
             # file I/O.
-            client.load_system_host_keys()
+            if self.host_key is not None:
+                # The pinned key replaces the system known-hosts file and
+                # applies to this host whatever names its line carries.
+                name = self.host if self.port == 22 else f"[{self.host}]:{self.port}"
+                client.get_host_keys().add(
+                    name, self.host_key.get_name(), self.host_key
+                )
+            else:
+                client.load_system_host_keys()
             if self.known_hosts is not None:
                 client.load_host_keys(str(self.known_hosts))
             client.set_missing_host_key_policy(
@@ -358,3 +410,13 @@ def _parse_df(df: str) -> tuple[float | None, float | None]:
         return round(int(parts[1]) / 1024, 2), float(parts[4].rstrip("%"))
     except (IndexError, ValueError):
         return None, None
+
+
+def _parse_host_key(known_hosts_line: str) -> paramiko.PKey:
+    try:
+        entry = HostKeyEntry.from_line(known_hosts_line)
+    except (paramiko.SSHException, InvalidHostKey) as e:
+        raise ValueError(f"Invalid host_key: {e}") from e
+    if entry is None:
+        raise ValueError("host_key must be a known-hosts line")
+    return entry.key

tests/test_ssh_client.py

@@ -3,13 +3,17 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Iterator
+import contextlib
 from datetime import UTC, datetime
 from pathlib import Path
+import socket
 import threading
 from typing import Any
 from unittest.mock import Mock, patch
 
 import paramiko
+from paramiko.common import AUTH_FAILED, AUTH_SUCCESSFUL
 import pytest
 
 from nanokvm.exceptions import NanoKVMNotSupportedError
@@ -20,6 +24,7 @@
     NanoKVMSSHConnectionError,
     NanoKVMSSHNotConnectedError,
     SSHMetrics,
+    probe_host_key,
 )
 
 
@@ -549,3 +554,106 @@ async def test_set_watchdog_enabled_rejects_firmware_without_watchdog() -> None:
         pytest.raises(NanoKVMNotSupportedError),
     ):
         await ssh.set_watchdog_enabled(True)
+
+
+class _PasswordServer(paramiko.ServerInterface):
+    def get_allowed_auths(self, username: str) -> str:
+        return "password"
+
+    def check_auth_password(self, username: str, password: str) -> int:
+        if password == "synthetic-password":
+            return AUTH_SUCCESSFUL
+        return AUTH_FAILED
+
+
+@pytest.fixture(scope="module")
+def server_key() -> paramiko.RSAKey:
+    return paramiko.RSAKey.generate(2048)
+
+
+@pytest.fixture
+def ssh_server(server_key: paramiko.RSAKey) -> Iterator[int]:
+    """Serve SSH handshakes and password logins on a local port."""
+    listener = socket.create_server(("127.0.0.1", 0))
+    transports: list[paramiko.Transport] = []
+
+    def serve() -> None:
+        while True:
+            try:
+                sock, _ = listener.accept()
+            except OSError:
+                return
+            transport = paramiko.Transport(sock)
+            transport.add_server_key(server_key)
+            transports.append(transport)
+            with contextlib.suppress(paramiko.SSHException, EOFError):
+                transport.start_server(server=_PasswordServer())
+
+    thread = threading.Thread(target=serve, daemon=True)
+    thread.start()
+    yield listener.getsockname()[1]
+    listener.close()
+    for transport in transports:
+        transport.close()
+    thread.join(1)
+
+
+@pytest.mark.asyncio
+async def test_probe_host_key_reads_key_without_logging_in(
+    ssh_server: int, server_key: paramiko.RSAKey
+) -> None:
+    host_key = await probe_host_key("127.0.0.1", ssh_server)
+
+    assert host_key.host == "127.0.0.1"
+    assert host_key.known_hosts_line == (
+        f"[127.0.0.1]:{ssh_server} ssh-rsa {server_key.get_base64()}"
+    )
+    assert host_key.fingerprint == server_key.fingerprint
+
+
+@pytest.mark.asyncio
+async def test_probe_host_key_reports_unreachable_host() -> None:
+    listener = socket.create_server(("127.0.0.1", 0))
+    port = listener.getsockname()[1]
+    listener.close()
+
+    with pytest.raises(NanoKVMSSHConnectionError):
+        await probe_host_key("127.0.0.1", port)
+
+
+@pytest.mark.asyncio
+async def test_pinned_host_key_authenticates_under_any_host_name(
+    ssh_server: int, server_key: paramiko.RSAKey
+) -> None:
+    line = f"old-name.local ssh-rsa {server_key.get_base64()}"
+
+    async with NanoKVMSSH("127.0.0.1", port=ssh_server, host_key=line) as ssh:
+        await ssh.authenticate("synthetic-password")
+        assert ssh.ssh_client is not None
+
+
+@pytest.mark.asyncio
+async def test_pinned_host_key_rejects_a_different_server_key(ssh_server: int) -> None:
+    other = paramiko.RSAKey.generate(2048)
+    line = f"127.0.0.1 ssh-rsa {other.get_base64()}"
+
+    async with NanoKVMSSH("127.0.0.1", port=ssh_server, host_key=line) as ssh:
+        with pytest.raises(NanoKVMSSHConnectionError):
+            await ssh.authenticate("synthetic-password")
+
+
+@pytest.mark.asyncio
+async def test_pinned_host_key_skips_system_known_hosts() -> None:
+    fake_client = Mock()
+    key = paramiko.RSAKey.generate(1024)
+    with patch("nanokvm.ssh_client.paramiko.SSHClient", return_value=fake_client):
+        client = NanoKVMSSH("kvm.local", host_key=f"x ssh-rsa {key.get_base64()}")
+        await client.authenticate("synthetic-password")
+
+    fake_client.load_system_host_keys.assert_not_called()
+
+
+@pytest.mark.parametrize("line", ["", "not a key line", "host ssh-rsa !!!"])
+def test_invalid_host_key_is_rejected_up_front(line: str) -> None:
+    with pytest.raises(ValueError):
+        NanoKVMSSH("kvm.local", host_key=line)