init

example/autoupdate.py

@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+"""Minimal self-update helper for Python scripts."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import stat
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any, Callable
+from urllib.error import HTTPError, URLError
+from urllib.parse import quote, urljoin
+from urllib.request import Request, urlopen
+
+
+class UpdateError(RuntimeError):
+    """Raised when the update flow cannot be completed safely."""
+
+
+ReleaseHandler = Callable[[Path, dict[str, Any]], None]
+
+
+def ensure_updated(
+    base_url: str,
+    api_key: str,
+    project_slug: str,
+    current_version: str,
+    apply_update: ReleaseHandler | None = None,
+) -> bool:
+    """
+    Check for a newer release and restart the current script after update.
+
+    If `apply_update` is omitted, the current script file is replaced in place.
+    If `apply_update` is provided, it receives `(downloaded_artifact_path, release_dict)`
+    and is responsible for applying the downloaded artifact before restart.
+    """
+    if not api_key or api_key == "upsk_replace_me":
+        raise UpdateError("set a real API key before calling ensure_updated")
+
+    script_path = _current_script_path()
+    release = _load_latest_release(
+        base_url=base_url,
+        api_key=api_key,
+        project_slug=project_slug,
+    )
+
+    latest_version = str(release.get("version") or "")
+    if not latest_version or latest_version == current_version:
+        return False
+
+    download_url = urljoin(base_url.rstrip("/") + "/", str(release.get("download_url") or ""))
+    if not download_url:
+        raise UpdateError("latest release does not contain a download URL")
+
+    temp_path = _download_release(
+        url=download_url,
+        api_key=api_key,
+        checksum_sha256=str(release.get("checksum_sha256") or ""),
+        size_bytes=int(release.get("size_bytes") or 0),
+        target_name=script_path.name,
+    )
+
+    try:
+        if apply_update is None:
+            _replace_file(temp_path, script_path)
+        else:
+            apply_update(temp_path, release)
+    finally:
+        temp_path.unlink(missing_ok=True)
+
+    _restart_current_process(script_path)
+    return True
+
+
+def _current_script_path() -> Path:
+    main_module = sys.modules.get("__main__")
+    main_file = getattr(main_module, "__file__", None)
+    if not main_file:
+        raise UpdateError("cannot resolve the current script path")
+    return Path(main_file).resolve()
+
+
+def _make_request(url: str, api_key: str, accept_json: bool = False) -> Request:
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "User-Agent": "autoupdate/1.0",
+    }
+    if accept_json:
+        headers["Accept"] = "application/json"
+    return Request(url, headers=headers)
+
+
+def _load_latest_release(
+    base_url: str,
+    api_key: str,
+    project_slug: str,
+) -> dict[str, Any]:
+    latest_url = urljoin(
+        base_url.rstrip("/") + "/",
+        f"api/v1/projects/{quote(project_slug, safe='')}/releases/latest",
+    )
+    request = _make_request(latest_url, api_key, accept_json=True)
+
+    try:
+        with urlopen(request, timeout=15) as response:
+            charset = response.headers.get_content_charset() or "utf-8"
+            payload = json.loads(response.read().decode(charset))
+    except HTTPError as exc:
+        if exc.code == 404:
+            raise UpdateError(f"no releases found for project '{project_slug}'") from exc
+        raise UpdateError(f"HTTP {exc.code} while checking updates: {exc.reason}") from exc
+    except URLError as exc:
+        raise UpdateError(f"network error while checking updates: {exc.reason}") from exc
+
+    release = payload.get("release")
+    if not isinstance(release, dict):
+        raise UpdateError("update server returned an unexpected response")
+    return release
+
+
+def _download_release(
+    url: str,
+    api_key: str,
+    checksum_sha256: str,
+    size_bytes: int,
+    target_name: str,
+) -> Path:
+    request = _make_request(url, api_key)
+    digest = hashlib.sha256()
+    written = 0
+    tmp_path: Path | None = None
+
+    try:
+        with urlopen(request, timeout=60) as response:
+            with tempfile.NamedTemporaryFile(
+                prefix=target_name + ".",
+                suffix=".download",
+                delete=False,
+            ) as tmp_file:
+                tmp_path = Path(tmp_file.name)
+                while True:
+                    chunk = response.read(64 * 1024)
+                    if not chunk:
+                        break
+                    tmp_file.write(chunk)
+                    digest.update(chunk)
+                    written += len(chunk)
+
+                tmp_file.flush()
+                os.fsync(tmp_file.fileno())
+    except HTTPError as exc:
+        if tmp_path is not None:
+            tmp_path.unlink(missing_ok=True)
+        raise UpdateError(f"HTTP {exc.code} while downloading update: {exc.reason}") from exc
+    except URLError as exc:
+        if tmp_path is not None:
+            tmp_path.unlink(missing_ok=True)
+        raise UpdateError(f"network error while downloading update: {exc.reason}") from exc
+    except Exception:
+        if tmp_path is not None:
+            tmp_path.unlink(missing_ok=True)
+        raise
+
+    actual_sha256 = digest.hexdigest()
+    if checksum_sha256 and actual_sha256 != checksum_sha256:
+        tmp_path.unlink(missing_ok=True)
+        raise UpdateError(
+            f"checksum mismatch: expected {checksum_sha256}, got {actual_sha256}"
+        )
+
+    if size_bytes and written != size_bytes:
+        tmp_path.unlink(missing_ok=True)
+        raise UpdateError(f"size mismatch: expected {size_bytes}, got {written}")
+
+    return tmp_path
+
+
+def _replace_file(downloaded_path: Path, target_path: Path) -> None:
+    target_path.parent.mkdir(parents=True, exist_ok=True)
+
+    if target_path.exists():
+        current_mode = stat.S_IMODE(target_path.stat().st_mode)
+        os.chmod(downloaded_path, current_mode)
+
+    os.replace(downloaded_path, target_path)
+
+
+def _restart_current_process(script_path: Path) -> None:
+    os.execv(sys.executable, [sys.executable, str(script_path), *sys.argv[1:]])