Implement an automatic update mechanism that has options to preserve …

docs/api.md

@@ -160,7 +160,7 @@ Autostart scripts run when the device boots.
 
 | Methods | Requires |
 |---|---|
-| `get_application_version`, `update_application` | all |
+| `get_application_version`, `update_application`, `install_application_update` | all, see below |
 | `get_preview_status`, `set_preview_state` | non-Pro 2.2.5 / Pro |
 | `update_application_offline` | non-Pro 2.3.1 |
 | `get_update_server`, `set_update_server` | non-Pro 2.5.1 |
@@ -170,8 +170,35 @@ Autostart scripts run when the device boots.
 | `assistant_install`, `assistant_start` | Pro 1.1.4 |
 | `kvmadmin_install`, `kvmadmin_uninstall`, `kvmadmin_start`, `kvmadmin_stop`, `kvmadmin_status` | Pro 1.1.5 |
 
+`install_application_update()` updates the application when the device is
+out of date, and waits until it runs the new version, which it returns. The
+update restarts the device's server and ends the session, so the client logs
+in again with credentials kept by `keep_credentials=True` (see
+[Connecting](connecting.md#keeping-credentials)):
+
+```python
+async with NanoKVMClient(url, keep_credentials=True) as client:
+    await client.authenticate("username", "password")
+    version = await client.install_application_update()
+```
+
+- Without kept credentials it raises `NanoKVMNotAuthenticatedError` before
+  contacting the device. It raises `NanoKVMError` when the device cannot reach
+  its update server, and `TimeoutError` when the new version is not running
+  within `timeout` (default 10 minutes).
+- `get_application_version()` returns the `current` and `latest` versions.
+  `latest` is empty when the device cannot reach its update server. The device
+  asks the server while answering, so the call waits up to 60 seconds.
+- `update_application()` answers once the update is installed, which can take
+  minutes; it waits up to 15 minutes. It does not wait for the restart. Non-Pro
+  firmware before 2.5.1 restarts before it answers, so the call returns when
+  the connection drops, without confirming the update; the same applies to
+  `update_application_offline()`. `install_application_update()` confirms it.
+- With `set_preview_state(True)`, the device installs preview releases.
+
 `update_application_offline("nanokvm_2.5.1.tar.gz")` installs an official
-package from a local file.
+package from a local file. Pass `sha256=` to check the file first; non-Pro
+2.5.1 and newer also check it on the device.
 
 ### Accounts
 

docs/connecting.md

@@ -39,6 +39,16 @@ plain text only after the device rejects the credentials. Pass
 `get_account()` returns the username and, when the firmware provides it, the
 role (older firmware returns `None`).
 
+### Keeping credentials
+
+The client does not keep the password. With `keep_credentials=True`, it keeps
+the username and password of a successful login in memory, so that it can log
+in again after the device restarts, as `install_application_update()` needs.
+Logging out, changing the password, a failed login and closing the client drop
+them. They are never written to disk or shown in `repr()`.
+
+### Sessions and logout
+
 Starting another login or logging out invalidates operations and streams still
 running on the old session; they raise `NanoKVMNotAuthenticatedError`. On
 current non-Pro firmware, logging out revokes **every** session on the device,

nanokvm/client.py

@@ -82,6 +82,7 @@ def __init__(
         ssl_ca_cert: str | None = None,
         ssl_fingerprint: str | None = None,
         use_password_obfuscation: bool | None = None,
+        keep_credentials: bool = False,
     ) -> None:
         """
         Initialize the NanoKVM client.
@@ -103,6 +104,9 @@ def __init__(
                 None = auto-detect (try obfuscated first, fall back to plain text).
                 True = always use obfuscated passwords (older NanoKVM versions).
                 False = always use plain text passwords (newer HTTPS-enabled versions).
+            keep_credentials: Keep the username and password of a successful
+                login in memory, to log in again after the device restarts, as
+                ``install_application_update()`` needs (default: False).
         """
         super().__init__(
             SessionController(
@@ -114,6 +118,7 @@ def __init__(
                 ssl_ca_cert=ssl_ca_cert,
                 ssl_fingerprint=ssl_fingerprint,
                 use_password_obfuscation=use_password_obfuscation,
+                keep_credentials=keep_credentials,
             )
         )
 

nanokvm/components/services.py

@@ -3,7 +3,9 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import hashlib
+import math
 from os import PathLike
 from pathlib import Path
 import re
@@ -16,6 +18,7 @@
     require_application_version,
     require_hardware,
 )
+from ..exceptions import NanoKVMError, NanoKVMNotAuthenticatedError
 from ..models.common import (
     GetPreviewRsp,
     GetTailscaleStatusRsp,
@@ -40,6 +43,10 @@
 
 _OFFLINE_UPDATE_PACKAGE_RE = re.compile(r"^nanokvm_[0-9]+\.[0-9]+\.[0-9]+\.tar\.gz$")
 _OFFLINE_UPDATE_TIMEOUT_SECONDS = 15 * 60
+# The device asks its update server while answering; the web UI waits 60 s.
+_VERSION_CHECK_TIMEOUT_SECONDS = 60.0
+# The device downloads and installs before answering; the web UI waits 15 min.
+_UPDATE_TIMEOUT_SECONDS = 15 * 60.0
 
 
 def _calculate_file_sha256(file_path: Path) -> str:
@@ -60,12 +67,18 @@ async def get_tailscale_status(self) -> GetTailscaleStatusRsp:
             response_model=GetTailscaleStatusRsp,
         )
 
-    async def get_application_version(self) -> GetVersionRsp:
-        """Get current and latest application versions."""
+    async def get_application_version(
+        self, *, timeout: float = _VERSION_CHECK_TIMEOUT_SECONDS
+    ) -> GetVersionRsp:
+        """Get current and latest application versions.
+
+        ``latest`` is empty when the device cannot reach its update server.
+        """
         return await self._session.api_request_json(
             hdrs.METH_GET,
             "/application/version",
             response_model=GetVersionRsp,
+            timeout=aiohttp.ClientTimeout(total=timeout),
         )
 
     @require_application_version(non_pro="2.2.5")
@@ -86,11 +99,114 @@ async def set_preview_state(self, enable: bool) -> None:
             data=SetPreviewReq(enable=enable),
         )
 
-    async def update_application(self) -> None:
-        """Trigger the application update process."""
-        await self._session.api_request_json(hdrs.METH_POST, "/application/update")
+    async def update_application(
+        self, *, timeout: float = _UPDATE_TIMEOUT_SECONDS
+    ) -> None:
+        """Install the latest application version.
+
+        The device answers once the update is installed, then restarts its
+        server, which ends the session. Non-Pro firmware before 2.5.1 restarts
+        before its answer is sent, so it drops the connection instead.
+        """
+        drops_connection = await self._restarts_before_answering()
+        try:
+            await self._session.api_request_json(
+                hdrs.METH_POST,
+                "/application/update",
+                timeout=aiohttp.ClientTimeout(total=timeout),
+            )
+        except aiohttp.ServerDisconnectedError:
+            if not drops_connection:
+                raise
+        except aiohttp.ClientResponseError as err:
+            if not self._is_restart_gateway_error(err):
+                raise
         self._session.forget_versions()
 
+    async def install_application_update(
+        self, *, timeout: float = 10 * 60, poll_interval: float = 2.0
+    ) -> str:
+        """Update the application if needed, and wait until the device runs it.
+
+        Returns the application version the device runs afterwards. The
+        restart ends the session, so the client logs in again with the
+        credentials it keeps with ``keep_credentials=True``. ``timeout`` bounds
+        the wait for the restart, in seconds.
+        """
+        for value in (timeout, poll_interval):
+            if not math.isfinite(value) or value <= 0:
+                raise ValueError("timeout and poll_interval must be positive seconds")
+        session = self._session
+        if session._credentials is None:
+            raise NanoKVMNotAuthenticatedError(
+                "install_application_update needs a client created with "
+                "keep_credentials=True that has authenticated"
+            )
+
+        versions = await self.get_application_version()
+        if not versions.latest:
+            raise NanoKVMError("The device could not reach its update server")
+        if _version_at_least(versions.current, versions.latest):
+            return versions.current
+
+        await self.update_application()
+        return await self._wait_for_new_application(
+            versions.current, timeout, poll_interval
+        )
+
+    async def _wait_for_new_application(
+        self, previous: str, timeout: float, poll_interval: float
+    ) -> str:
+        """Wait until the device answers with another application version."""
+        session = self._session
+        try:
+            async with asyncio.timeout(timeout):
+                while True:
+                    await asyncio.sleep(poll_interval)
+                    try:
+                        await session.detect_versions()
+                    except NanoKVMNotAuthenticatedError:
+                        # The restarted server rejects the old session. Log in
+                        # again; a server still starting is retried next time.
+                        with contextlib.suppress(aiohttp.ClientError, TimeoutError):
+                            await session.reauthenticate()
+                        continue
+                    except (aiohttp.ClientError, TimeoutError):
+                        continue  # Restarting.
+                    version = session._application_version
+                    if version is not None and version != previous:
+                        return version
+        except TimeoutError:
+            raise TimeoutError(
+                f"The device did not run a new application version within "
+                f"{timeout:g} seconds"
+            ) from None
+
+    async def _restarts_before_answering(self) -> bool:
+        """Whether an update restarts the server before its answer is sent.
+
+        Before 2.5.1, the non-Pro update handler runs the restart itself, which
+        stops the server while the answer is still buffered.
+        """
+        session = self._session
+        if await session.application_version_at_least(non_pro="2.5.1"):
+            return False
+        return session.is_hardware_family(HWFamily.NON_PRO)
+
+    def _is_restart_gateway_error(self, err: aiohttp.ClientResponseError) -> bool:
+        """Whether a 502 came from the device restarting after an update.
+
+        On non-Pro 2.5.1 and newer, the device can stop its old HTTP service
+        before the proxy has forwarded the successful response. Its official
+        UI treats that gateway response as successful too.
+        """
+        version = self._session._application_version
+        return (
+            err.status == 502
+            and version is not None
+            and _version_at_least(version, "2.5.1")
+        )
+
     @require_hardware(HWFamily.NON_PRO, min_version="2.3.1")
     async def update_application_offline(
         self,
@@ -122,22 +238,19 @@ async def update_application_offline(
             if checksum and supports_remote_checksum
             else {}
         )
+        drops_connection = await self._restarts_before_answering()
         try:
             await self._session.upload_file(
                 "/application/update/offline",
                 package_path,
                 headers=headers,
                 timeout=aiohttp.ClientTimeout(total=_OFFLINE_UPDATE_TIMEOUT_SECONDS),
             )
+        except aiohttp.ServerDisconnectedError:
+            if not drops_connection:
+                raise
         except aiohttp.ClientResponseError as err:
-            # The device can stop its old HTTP service before the proxy has
-            # forwarded the successful response. Its official UI treats that
-            # restart-related gateway response as successful too.
-            if (
-                err.status != 502
-                or self._session._application_version is None
-                or not _version_at_least(self._session._application_version, "2.5.1")
-            ):
+            if not self._is_restart_gateway_error(err):
                 raise
         self._session.forget_versions()
 

nanokvm/components/session.py

@@ -5,6 +5,7 @@
 import asyncio
 from collections.abc import AsyncGenerator
 import contextlib
+from dataclasses import dataclass, field
 from http.cookies import Morsel, SimpleCookie
 import json
 import logging
@@ -61,6 +62,14 @@
 _CURRENT_PASSWORD_MIN_NON_PRO_VERSION = (2, 5, 1)
 
 
+@dataclass(frozen=True)
+class _Credentials:
+    """Login credentials kept in memory to log in again."""
+
+    username: str
+    password: str = field(repr=False)
+
+
 class SessionController:
     """Own authentication, shared state, HTTP and WebSocket resources."""
 
@@ -75,6 +84,7 @@ def __init__(
         ssl_ca_cert: str | None = None,
         ssl_fingerprint: str | None = None,
         use_password_obfuscation: bool | None = None,
+        keep_credentials: bool = False,
     ) -> None:
         """Initialize the shared session, transport and device state."""
         self.url = yarl.URL(url)
@@ -94,6 +104,8 @@ def __init__(
         self._ssl_ca_cert = ssl_ca_cert
         self._ssl_fingerprint = ssl_fingerprint
         self._use_password_obfuscation = use_password_obfuscation
+        self._keep_credentials = keep_credentials
+        self._credentials: _Credentials | None = None
         self._ssl_config: ssl.SSLContext | Fingerprint | bool | None = None
         self._hw_version: HWVersion | None = None
         self._application_version: str | None = None
@@ -127,6 +139,7 @@ async def enter(self) -> None:
 
     async def exit(self) -> None:
         """Close the client's WebSocket and owned HTTP resources."""
+        self._credentials = None
         await self.close_ws()
 
         if self._ws_session is not None:
@@ -377,7 +390,35 @@ async def do_authenticate(
             raise
 
     async def authenticate(self, username: str, password: str) -> int:
-        """Authenticate and store the session token."""
+        """Authenticate and store the session token.
+
+        With ``keep_credentials``, a successful login also keeps the
+        credentials for ``reauthenticate()``.
+        """
+        self._credentials = None
+        generation = await self._login(username, password)
+        if self._keep_credentials:
+            self._credentials = _Credentials(username, password)
+        return generation
+
+    async def reauthenticate(self) -> int:
+        """Log in again with the credentials kept from the last login.
+
+        The credentials stay kept when the device cannot be reached, and are
+        dropped when it rejects them.
+        """
+        credentials = self._credentials
+        if credentials is None:
+            raise NanoKVMNotAuthenticatedError("No credentials are kept to log in")
+        try:
+            return await self._login(credentials.username, credentials.password)
+        except NanoKVMAuthenticationFailure:
+            if self._credentials is credentials:
+                self._credentials = None
+            raise
+
+    async def _login(self, username: str, password: str) -> int:
+        """Replace the session with a new login."""
         # A failed identity switch must never leave the previous account usable.
         generation = await self.clear_local_session()
         _LOGGER.debug("Attempting authentication for user: %s", username)
@@ -417,6 +458,7 @@ def check_session_generation(self, generation: int) -> None:
 
     async def logout(self) -> None:
         """Log out and clear the session token."""
+        self._credentials = None
         generation = self._session_generation
         try:
             if self._token and self._token != "disabled":
@@ -505,8 +547,12 @@ async def change_password(
         *,
         current_password: str | None = None,
     ) -> None:
-        """Change the KVM password for the authenticated account."""
+        """Change the KVM password for the authenticated account.
+
+        Kept credentials are dropped; authenticate again to keep the new ones.
+        """
         generation = self._session_generation
+        self._credentials = None
         if await self.uses_current_password_contract():
             if not current_password or not current_password.strip():
                 raise ValueError(