277 lines
8.5 KiB
Python
277 lines
8.5 KiB
Python
#!/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]
|
|
DEFAULT_CONFIG_NAME = "autoupdate.json"
|
|
DEFAULT_CONFIG = {
|
|
"base_url": "https://updates.example.com",
|
|
"api_key": "upsk_replace_me",
|
|
"project_slug": "example-project",
|
|
}
|
|
|
|
|
|
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 _is_placeholder_config(
|
|
{
|
|
"base_url": base_url,
|
|
"api_key": api_key,
|
|
"project_slug": project_slug,
|
|
}
|
|
):
|
|
print("autoupdate: configure autoupdate.json to enable updates")
|
|
return False
|
|
|
|
if not api_key:
|
|
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
|
|
|
|
print(f"autoupdate: updating {current_version} -> {latest_version}")
|
|
|
|
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 ensure_updated_from_config(
|
|
current_version: str,
|
|
config_path: str | Path | None = None,
|
|
apply_update: ReleaseHandler | None = None,
|
|
) -> bool:
|
|
"""
|
|
Load update settings from JSON and run the update check.
|
|
|
|
If `config_path` is omitted, `autoupdate.json` is loaded from the same
|
|
directory as the currently running script.
|
|
"""
|
|
config = _load_config(config_path)
|
|
return ensure_updated(
|
|
base_url=config["base_url"],
|
|
api_key=config["api_key"],
|
|
project_slug=config["project_slug"],
|
|
current_version=current_version,
|
|
apply_update=apply_update,
|
|
)
|
|
|
|
|
|
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 _default_config_path() -> Path:
|
|
return _current_script_path().with_name(DEFAULT_CONFIG_NAME)
|
|
|
|
|
|
def _load_config(config_path: str | Path | None) -> dict[str, str]:
|
|
path = (
|
|
Path(config_path).expanduser()
|
|
if config_path is not None
|
|
else _default_config_path()
|
|
)
|
|
try:
|
|
with path.open("r", encoding="utf-8") as config_file:
|
|
payload = json.load(config_file)
|
|
except FileNotFoundError:
|
|
_write_default_config(path)
|
|
return DEFAULT_CONFIG.copy()
|
|
except json.JSONDecodeError as exc:
|
|
raise UpdateError(f"update config file is not valid JSON: {path}") from exc
|
|
|
|
if not isinstance(payload, dict):
|
|
raise UpdateError(f"update config must be a JSON object: {path}")
|
|
|
|
config: dict[str, str] = {}
|
|
for key in ("base_url", "api_key", "project_slug"):
|
|
value = payload.get(key)
|
|
if not isinstance(value, str) or not value.strip():
|
|
raise UpdateError(f"update config field {key!r} is required: {path}")
|
|
config[key] = value.strip()
|
|
|
|
return config
|
|
|
|
|
|
def _write_default_config(path: Path) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("w", encoding="utf-8") as config_file:
|
|
json.dump(DEFAULT_CONFIG, config_file, indent=2)
|
|
config_file.write("\n")
|
|
|
|
|
|
def _is_placeholder_config(config: dict[str, str]) -> bool:
|
|
return any(
|
|
config.get(key, "").strip() == value
|
|
for key, value in DEFAULT_CONFIG.items()
|
|
)
|
|
|
|
|
|
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:]])
|