pristontale-eu-linux/ptlinux/download.py

"""Secure, conditional download of the official PristonTale EU launcher."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from datetime import datetime
import hashlib
import json
import os
from pathlib import Path
import struct
from typing import Callable
from urllib.error import HTTPError
from urllib.request import Request, urlopen

from . import __version__


OFFICIAL_LAUNCHER_URL = (
	"https://pristontale.eu/storage/uploads/Pristontale%20EU%20Launcher.exe"
)
MAX_LAUNCHER_BYTES = 64 * 1024 * 1024
READ_SIZE = 256 * 1024
USER_AGENT = f"PT-Linux/{__version__} (+https://pristontale.eu/)"


class DownloadError(RuntimeError):
	"""Raised when the official launcher cannot be safely downloaded."""


@dataclass(frozen=True)
class LauncherDownload:
	path: str
	url: str
	sha256: str
	size: int
	etag: str | None
	last_modified: str | None
	fetched_at: str
	changed: bool


def fetch_official_launcher(
	destination: Path,
	metadata_path: Path,
	*,
	url: str = OFFICIAL_LAUNCHER_URL,
	timeout_seconds: float = 60.0,
	progress: Callable[[int, int | None], None] | None = None,
) -> LauncherDownload:
	"""Download the current official launcher and atomically replace the cache."""

	metadata = _read_metadata(metadata_path)
	headers = {"User-Agent": USER_AGENT}
	if destination.is_file() and _looks_like_pe(destination):
		etag = metadata.get("etag")
		last_modified = metadata.get("last_modified")
		if isinstance(etag, str) and etag:
			headers["If-None-Match"] = etag
		if isinstance(last_modified, str) and last_modified:
			headers["If-Modified-Since"] = last_modified

	request = Request(url, headers=headers, method="GET")
	try:
		response = urlopen(request, timeout=timeout_seconds)
	except HTTPError as exc:
		status = exc.code
		exc.close()
		if status == 304 and destination.is_file():
			return _existing_download(destination, metadata, url)
		raise DownloadError(f"official launcher request failed: HTTP {status}") from exc
	except OSError as exc:
		raise DownloadError(f"official launcher request failed: {exc}") from exc

	with response:
		length = _content_length(response.headers.get("Content-Length"))
		if length is not None and length > MAX_LAUNCHER_BYTES:
			raise DownloadError(
				f"launcher response is unexpectedly large: {length} bytes"
			)
		destination.parent.mkdir(parents=True, exist_ok=True)
		temporary = destination.with_name(f".{destination.name}.{os.getpid()}.part")
		digest = hashlib.sha256()
		written = 0
		try:
			with temporary.open("wb") as handle:
				while True:
					chunk = response.read(READ_SIZE)
					if not chunk:
						break
					written += len(chunk)
					if written > MAX_LAUNCHER_BYTES:
						raise DownloadError(
							"launcher response exceeded the 64 MiB safety limit"
						)
					digest.update(chunk)
					handle.write(chunk)
					if progress is not None:
						progress(written, length)
				handle.flush()
				os.fsync(handle.fileno())
			if length is not None and written != length:
				raise DownloadError(
					f"incomplete launcher response: {written}/{length} bytes"
				)
			validate_pe_launcher(temporary)
			temporary.replace(destination)
		except BaseException:
			temporary.unlink(missing_ok=True)
			raise

	result = LauncherDownload(
		path=str(destination),
		url=url,
		sha256=digest.hexdigest(),
		size=written,
		etag=response.headers.get("ETag"),
		last_modified=response.headers.get("Last-Modified"),
		fetched_at=datetime.now().astimezone().isoformat(),
		changed=True,
	)
	_write_json(metadata_path, asdict(result))
	return result


def validate_pe_launcher(path: Path) -> None:
	"""Validate the minimum PE structure and require a 64-bit GUI executable."""

	try:
		with path.open("rb") as handle:
			dos = handle.read(64)
			if len(dos) < 64 or dos[:2] != b"MZ":
				raise DownloadError("download is not a DOS/PE executable")
			pe_offset = struct.unpack_from("<I", dos, 0x3C)[0]
			if pe_offset < 64 or pe_offset > 8 * 1024 * 1024:
				raise DownloadError("launcher has an invalid PE header offset")
			handle.seek(pe_offset)
			pe_header = handle.read(24)
			if len(pe_header) != 24 or pe_header[:4] != b"PE\x00\x00":
				raise DownloadError("launcher has no valid PE signature")
			machine = struct.unpack_from("<H", pe_header, 4)[0]
			if machine != 0x8664:
				raise DownloadError(
					f"launcher has unexpected PE machine type 0x{machine:04x}"
				)
	except OSError as exc:
		raise DownloadError(f"cannot validate launcher: {exc}") from exc


def _looks_like_pe(path: Path) -> bool:
	try:
		validate_pe_launcher(path)
	except DownloadError:
		return False
	return True


def _existing_download(
	path: Path,
	metadata: dict[str, object],
	url: str,
) -> LauncherDownload:
	validate_pe_launcher(path)
	return LauncherDownload(
		path=str(path),
		url=url,
		sha256=_sha256(path),
		size=path.stat().st_size,
		etag=_optional_string(metadata.get("etag")),
		last_modified=_optional_string(metadata.get("last_modified")),
		fetched_at=datetime.now().astimezone().isoformat(),
		changed=False,
	)


def _read_metadata(path: Path) -> dict[str, object]:
	try:
		value = json.loads(path.read_text(encoding="utf-8"))
	except (OSError, json.JSONDecodeError):
		return {}
	return value if isinstance(value, dict) else {}


def _write_json(path: Path, value: object) -> None:
	path.parent.mkdir(parents=True, exist_ok=True)
	temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
	temporary.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8")
	temporary.replace(path)


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 _content_length(value: str | None) -> int | None:
	if value is None:
		return None
	try:
		length = int(value)
	except ValueError as exc:
		raise DownloadError(f"invalid Content-Length: {value!r}") from exc
	if length < 0:
		raise DownloadError(f"invalid Content-Length: {value!r}")
	return length


def _optional_string(value: object) -> str | None:
	return value if isinstance(value, str) else None