pristontale-eu-linux/ptlinux/umu.py

"""Pinned, user-level bootstrap for the official UMU launcher zipapp."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from datetime import datetime
import hashlib
import json
import os
from pathlib import Path
import secrets
import shutil
import sys
import tarfile
import threading
import time
from typing import Callable, Mapping, Protocol
from urllib.error import HTTPError
from urllib.request import Request, urlopen
import zipfile

from . import __version__
from .paths import AppPaths
from .telemetry import (
	classify_network_failure,
	current_telemetry,
	network_http_status,
)


UMU_LAUNCHER_VERSION = "1.4.4"
UMU_LAUNCHER_URL = (
	"https://github.com/Open-Wine-Components/umu-launcher/releases/download/"
	f"{UMU_LAUNCHER_VERSION}/umu-launcher-{UMU_LAUNCHER_VERSION}-zipapp.tar"
)
UMU_ARCHIVE_SHA256 = "eb590691841f7fad3fc3ad8fd5db4ccb87849fe7948e62b28ece7a4ee48cc851"
UMU_EXECUTABLE_SHA256 = "d0005a58602041229cc467dab03dc0c0b9e8cce09a8145b16b7683244cf17804"
UMU_ARCHIVE_MEMBER = "umu/umu-run"
MAX_UMU_ARCHIVE_BYTES = 4 * 1024 * 1024
MAX_UMU_EXECUTABLE_BYTES = 2 * 1024 * 1024
READ_SIZE = 256 * 1024
USER_AGENT = f"PT-Linux/{__version__} (+https://pristontale.eu/)"


class UmuBootstrapError(RuntimeError):
	"""Raised when the pinned portable UMU launcher cannot be prepared."""


class _HttpResponse(Protocol):
	headers: Mapping[str, str]

	def read(self, size: int = -1) -> bytes:
		...

	def __enter__(self) -> _HttpResponse:
		...

	def __exit__(
		self,
		exception_type: object,
		exception: object,
		traceback: object,
	) -> object:
		...


UmuOpener = Callable[..., _HttpResponse]


@dataclass(frozen=True)
class UmuBootstrapResult:
	path: str
	version: str
	source_url: str
	archive_sha256: str
	executable_sha256: str
	changed: bool


def find_umu_launcher(
	paths: AppPaths | None = None,
	environment: Mapping[str, str] | None = None,
) -> str | None:
	"""Find a system UMU first, then the launcher's verified portable copy."""

	env = os.environ if environment is None else environment
	located = shutil.which("umu-run", path=env.get("PATH"))
	if located:
		candidate = Path(located).expanduser().resolve()
		if (
			paths is None
			or candidate != paths.managed_umu_executable.resolve()
			or managed_umu_ready(candidate)
		):
			return str(candidate)
	if paths is not None and managed_umu_ready(paths.managed_umu_executable):
		return str(paths.managed_umu_executable.resolve())
	return None


def managed_umu_ready(path: Path) -> bool:
	"""Accept only the exact launcher zipapp pinned by this release."""

	try:
		return (
			path.is_file()
			and not path.is_symlink()
			and os.access(path, os.X_OK)
			and _sha256(path) == UMU_EXECUTABLE_SHA256
		)
	except OSError:
		return False


def provision_umu(
	paths: AppPaths,
	*,
	environment: Mapping[str, str] | None = None,
	timeout_seconds: float = 60.0,
	opener: UmuOpener = urlopen,
) -> UmuBootstrapResult | None:
	"""Ensure that either a system or launcher-managed UMU is available."""

	if find_umu_launcher(paths, environment) is not None:
		return None
	return ensure_managed_umu(
		paths,
		timeout_seconds=timeout_seconds,
		opener=opener,
	)


def provision_umu_for_startup(
	paths: AppPaths,
	*,
	timeout_seconds: float = 15.0,
) -> UmuBootstrapResult | None:
	"""Repair a missing helper without preventing dashboard startup on failure."""

	try:
		return provision_umu(paths, timeout_seconds=timeout_seconds)
	except (OSError, RuntimeError) as exc:
		print(
			"Warning: the Steam Runtime helper could not be prepared during launcher "
			f"startup; Play will retry ({exc}).",
			file=sys.stderr,
		)
		return None


def ensure_managed_umu(
	paths: AppPaths,
	*,
	timeout_seconds: float = 60.0,
	opener: UmuOpener = urlopen,
) -> UmuBootstrapResult:
	"""Install the pinned official zipapp without modifying the host Python."""

	destination = paths.managed_umu_executable
	if managed_umu_ready(destination):
		return _result(destination, changed=False)

	archive = paths.managed_umu_archive(UMU_LAUNCHER_VERSION)
	if not _verified_archive(archive):
		telemetry = current_telemetry()
		download_started_at = time.monotonic()
		if telemetry is not None:
			telemetry.send_event(
				"download_started",
				metadata={
					"resource_type": "runtime",
					"mirror_id": "github.com",
				},
				context={"operation": "runtime_install", "stage": "download"},
			)
		try:
			_download_archive(
				archive,
				timeout_seconds=timeout_seconds,
				opener=opener,
			)
		except BaseException as exc:
			if telemetry is not None:
				telemetry.send_event(
					"download_failed",
					severity="error",
					message=str(exc),
					metadata={
						"resource_type": "runtime",
						"duration_ms": round(
							(time.monotonic() - download_started_at) * 1000
						),
						"failure_type": classify_network_failure(exc),
						"http_status": network_http_status(exc),
						"retry_count": 0,
						"resume_used": False,
						"mirror_id": "github.com",
					},
					context={"operation": "runtime_install", "stage": "download"},
				)
			raise
		download_duration = max(0.0, time.monotonic() - download_started_at)
		download_bytes = archive.stat().st_size
		if telemetry is not None:
			telemetry.send_event(
				"download_completed",
				metadata={
					"resource_type": "runtime",
					"download_bytes": download_bytes,
					"content_bytes": download_bytes,
					"duration_ms": round(download_duration * 1000),
					"average_mbps": _average_mbps(
						download_bytes,
						download_duration,
					),
					"retry_count": 0,
					"resume_used": False,
					"http_status": 200,
					"mirror_id": "github.com",
				},
				context={"operation": "runtime_install", "stage": "download"},
			)
	_install_zipapp(archive, destination)
	result = _result(destination, changed=True)
	_write_json(
		paths.managed_umu_metadata,
		{
			**asdict(result),
			"installed_at": datetime.now().astimezone().isoformat(),
		},
	)
	return result


def _download_archive(
	destination: Path,
	*,
	timeout_seconds: float,
	opener: UmuOpener,
) -> None:
	request = Request(
		UMU_LAUNCHER_URL,
		headers={"User-Agent": USER_AGENT},
		method="GET",
	)
	try:
		response = opener(request, timeout=timeout_seconds)
	except HTTPError as exc:
		status = exc.code
		exc.close()
		raise UmuBootstrapError(
			f"portable UMU request failed: HTTP {status}"
		) from exc
	except OSError as exc:
		raise UmuBootstrapError(f"portable UMU request failed: {exc}") from exc

	with response:
		length = _content_length(response.headers.get("Content-Length"))
		if length is not None and length > MAX_UMU_ARCHIVE_BYTES:
			raise UmuBootstrapError(
				f"portable UMU response is unexpectedly large: {length} bytes"
			)
		destination.parent.mkdir(parents=True, exist_ok=True)
		temporary = _temporary_path(destination, "part")
		digest = hashlib.sha256()
		written = 0
		try:
			with temporary.open("xb") as handle:
				while True:
					chunk = response.read(READ_SIZE)
					if not chunk:
						break
					written += len(chunk)
					if written > MAX_UMU_ARCHIVE_BYTES:
						raise UmuBootstrapError(
							"portable UMU response exceeded the 4 MiB safety limit"
						)
					digest.update(chunk)
					handle.write(chunk)
				handle.flush()
				os.fsync(handle.fileno())
			if length is not None and written != length:
				raise UmuBootstrapError(
					f"incomplete portable UMU response: {written}/{length} bytes"
				)
			actual_hash = digest.hexdigest()
			if actual_hash != UMU_ARCHIVE_SHA256:
				raise UmuBootstrapError(
					"portable UMU archive failed SHA-256 verification: "
					f"expected {UMU_ARCHIVE_SHA256}, received {actual_hash}"
				)
			temporary.replace(destination)
		except BaseException:
			temporary.unlink(missing_ok=True)
			raise


def _average_mbps(download_bytes: int, duration_seconds: float) -> float | None:
	if download_bytes <= 0 or duration_seconds <= 0:
		return None
	return round((download_bytes * 8) / duration_seconds / 1_000_000, 3)


def _install_zipapp(archive: Path, destination: Path) -> None:
	destination.parent.mkdir(parents=True, exist_ok=True)
	temporary = _temporary_path(destination, "tmp")
	try:
		try:
			with tarfile.open(archive, "r:") as bundle:
				member = bundle.getmember(UMU_ARCHIVE_MEMBER)
				if not member.isfile():
					raise UmuBootstrapError(
						"portable UMU archive member is not a regular file"
					)
				if member.size <= 0 or member.size > MAX_UMU_EXECUTABLE_BYTES:
					raise UmuBootstrapError(
						f"portable UMU executable has invalid size: {member.size} bytes"
					)
				source = bundle.extractfile(member)
				if source is None:
					raise UmuBootstrapError(
						"portable UMU executable could not be read from its archive"
					)
				with source, temporary.open("xb") as handle:
					digest = hashlib.sha256()
					written = 0
					while True:
						chunk = source.read(READ_SIZE)
						if not chunk:
							break
						written += len(chunk)
						digest.update(chunk)
						handle.write(chunk)
					handle.flush()
					os.fsync(handle.fileno())
		except (KeyError, OSError, tarfile.TarError) as exc:
			raise UmuBootstrapError(f"invalid portable UMU archive: {exc}") from exc
		if written != member.size or digest.hexdigest() != UMU_EXECUTABLE_SHA256:
			raise UmuBootstrapError(
				"portable UMU executable failed content verification"
			)
		if not zipfile.is_zipfile(temporary):
			raise UmuBootstrapError("portable UMU executable is not a valid zipapp")
		with zipfile.ZipFile(temporary) as application:
			if "__main__.py" not in application.namelist():
				raise UmuBootstrapError(
					"portable UMU zipapp has no __main__.py entry point"
				)
			bad_member = application.testzip()
			if bad_member is not None:
				raise UmuBootstrapError(
					f"portable UMU zipapp contains a corrupt member: {bad_member}"
				)
		temporary.chmod(0o755)
		temporary.replace(destination)
	except BaseException:
		temporary.unlink(missing_ok=True)
		raise


def _verified_archive(path: Path) -> bool:
	try:
		return (
			path.is_file()
			and not path.is_symlink()
			and path.stat().st_size <= MAX_UMU_ARCHIVE_BYTES
			and _sha256(path) == UMU_ARCHIVE_SHA256
		)
	except OSError:
		return False


def _result(path: Path, *, changed: bool) -> UmuBootstrapResult:
	return UmuBootstrapResult(
		path=str(path),
		version=UMU_LAUNCHER_VERSION,
		source_url=UMU_LAUNCHER_URL,
		archive_sha256=UMU_ARCHIVE_SHA256,
		executable_sha256=UMU_EXECUTABLE_SHA256,
		changed=changed,
	)


def _content_length(value: str | None) -> int | None:
	if value is None:
		return None
	try:
		length = int(value)
	except ValueError as exc:
		raise UmuBootstrapError(f"invalid portable UMU Content-Length: {value!r}") from exc
	if length < 0:
		raise UmuBootstrapError(f"invalid portable UMU Content-Length: {value!r}")
	return length


def _sha256(path: Path) -> str:
	digest = hashlib.sha256()
	with path.open("rb") as handle:
		for chunk in iter(lambda: handle.read(READ_SIZE), b""):
			digest.update(chunk)
	return digest.hexdigest()


def _temporary_path(path: Path, suffix: str) -> Path:
	return path.with_name(
		f".{path.name}.{os.getpid()}.{threading.get_ident()}."
		f"{secrets.token_hex(8)}.{suffix}"
	)


def _write_json(path: Path, value: object) -> None:
	path.parent.mkdir(parents=True, exist_ok=True)
	temporary = _temporary_path(path, "tmp")
	try:
		with temporary.open("x", encoding="utf-8", newline="\n") as handle:
			json.dump(value, handle, indent=2, sort_keys=True)
			handle.write("\n")
			handle.flush()
			os.fsync(handle.fileno())
		temporary.replace(path)
	except BaseException:
		temporary.unlink(missing_ok=True)
		raise