pristontale-eu-linux/src/ptlinux/importer.py

"""Safe local bootstrap from an existing PristonTale EU installation."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from datetime import datetime
import hashlib
import json
import os
from pathlib import Path
import shutil
import stat
from typing import Mapping
import uuid

from .command import CommandRunner
from .install_status import REQUIRED_GAME_FILES
from .overlay import COMPATIBILITY_BINARY_HASHES, cache_compatibility_binaries
from .paths import AppPaths
from .settings import Settings


class GameImportError(RuntimeError):
	"""Raised when a source cannot be safely imported."""


@dataclass(frozen=True)
class GameTreeInventory:
	files: int
	bytes: int
	core_sha256: dict[str, str]


@dataclass(frozen=True)
class GameImportResult:
	source: str
	destination: str
	copy_mode: str
	files: int
	bytes: int
	core_sha256: dict[str, str]
	depot_metadata_source: str | None
	depot_metadata_destination: str | None
	installed_build: int | None
	started_at: str
	completed_at: str
	verification_required: bool


def import_existing_game(
	paths: AppPaths,
	source_game: Path,
	runner: CommandRunner,
	*,
	environment: Mapping[str, str] | None = None,
) -> GameImportResult:
	"""Clone a complete source without ever writing through to it."""

	source = source_game.expanduser().resolve(strict=True)
	destination = paths.source_game
	_validate_destination(paths, destination)
	if source == destination.resolve(strict=False):
		raise GameImportError("source and managed destination are the same directory")
	if destination.resolve(strict=False).is_relative_to(source):
		raise GameImportError("managed destination cannot be inside the source game tree")
	if destination.exists() or destination.is_symlink():
		raise GameImportError(
			f"managed game directory already exists: {destination}; "
			"finish or remove the existing updater installation first"
		)

	started_at = datetime.now().astimezone().isoformat()
	source_before = inventory_game_tree(source)
	depot_metadata = discover_depot_metadata(source)
	if runner.dry_run:
		return GameImportResult(
			source=str(source),
			destination=str(destination),
			copy_mode="reflink-or-copy-dry-run",
			files=source_before.files,
			bytes=source_before.bytes,
			core_sha256=source_before.core_sha256,
			depot_metadata_source=(
				str(depot_metadata.path) if depot_metadata is not None else None
			),
			depot_metadata_destination=None,
			installed_build=(
				depot_metadata.installed_build if depot_metadata is not None else None
			),
			started_at=started_at,
			completed_at=started_at,
			verification_required=True,
		)

	destination.parent.mkdir(parents=True, exist_ok=True)
	temporary = destination.with_name(f".{destination.name}.import-{uuid.uuid4().hex}")
	committed = False
	pending_recorded = False
	try:
		temporary.mkdir(parents=False, exist_ok=False)
		runner.run(
			["cp", "-a", "--reflink=auto", "--", f"{source}/.", str(temporary)],
			timeout_seconds=1800.0,
			telemetry_stage="copy",
		)
		source_after = inventory_game_tree(source)
		copied = inventory_game_tree(temporary)
		if source_after != source_before:
			raise GameImportError("source game changed during import; staged copy discarded")
		if copied != source_before:
			raise GameImportError(
				"staged game inventory does not match the source; staged copy discarded"
			)
		if all(
			source_before.core_sha256.get(name) == expected
			for name, expected in COMPATIBILITY_BINARY_HASHES.items()
		):
			cache_compatibility_binaries(paths, source)
		temporary.replace(destination)
		committed = True
		provisional = GameImportResult(
			source=str(source),
			destination=str(destination),
			copy_mode="reflink-or-copy",
			files=source_before.files,
			bytes=source_before.bytes,
			core_sha256=source_before.core_sha256,
			depot_metadata_source=(
				str(depot_metadata.path) if depot_metadata is not None else None
			),
			depot_metadata_destination=None,
			installed_build=(
				depot_metadata.installed_build if depot_metadata is not None else None
			),
			started_at=started_at,
			completed_at=datetime.now().astimezone().isoformat(),
			verification_required=True,
		)
		_write_import_manifest(paths, provisional)
		pending_recorded = True
	except BaseException:
		_remove_staging_tree(temporary, destination.parent)
		if committed and not pending_recorded:
			_remove_staging_tree(destination, destination.parent)
		raise

	metadata_destination = None
	if depot_metadata is not None:
		metadata_destination = _install_depot_metadata(
			paths,
			depot_metadata,
			environment,
		)
	result = GameImportResult(
		source=str(source),
		destination=str(destination),
		copy_mode="reflink-or-copy",
		files=source_before.files,
		bytes=source_before.bytes,
		core_sha256=source_before.core_sha256,
		depot_metadata_source=(
			str(depot_metadata.path) if depot_metadata is not None else None
		),
		depot_metadata_destination=(
			str(metadata_destination) if metadata_destination is not None else None
		),
		installed_build=(
			depot_metadata.installed_build if depot_metadata is not None else None
		),
		started_at=started_at,
		completed_at=datetime.now().astimezone().isoformat(),
		verification_required=True,
	)
	_write_import_manifest(paths, result)
	return result


@dataclass(frozen=True)
class DepotMetadata:
	path: Path
	installed_build: int
	sha256: str


def discover_depot_metadata(source_game: Path) -> DepotMetadata | None:
	source = source_game.expanduser().resolve(strict=True)
	prefix = None
	for parent in source.parents:
		if parent.name.casefold() == "drive_c":
			prefix = parent.parent
			break
	if prefix is None:
		return None
	users = prefix / "drive_c/users"
	try:
		user_directories = sorted(path for path in users.iterdir() if path.is_dir())
	except OSError:
		return None
	for user_directory in user_directories:
		candidate = (
			user_directory
			/ "AppData/Roaming/Pristontale EU Launcher/Depots/depot_1.json"
		)
		try:
			value = json.loads(candidate.read_text(encoding="utf-8-sig"))
		except (OSError, json.JSONDecodeError):
			continue
		if not isinstance(value, dict):
			continue
		installed_build = value.get("InstalledBuild")
		if (
			isinstance(installed_build, bool)
			or not isinstance(installed_build, int)
			or installed_build <= 0
			or value.get("UpdateComplete") is not True
		):
			continue
		install_folder = value.get("InstallFolder")
		if not isinstance(install_folder, str) or not install_folder.casefold().endswith(
			r"\program files (x86)\pristontale eu".casefold()
		):
			continue
		return DepotMetadata(
			path=candidate,
			installed_build=installed_build,
			sha256=_sha256(candidate),
		)
	return None


def inventory_game_tree(source_game: Path) -> GameTreeInventory:
	source = source_game.expanduser().resolve(strict=True)
	for relative in REQUIRED_GAME_FILES:
		path = source / relative
		try:
			status = path.lstat()
		except OSError:
			raise GameImportError(f"game installation is incomplete: missing {relative}")
		if stat.S_ISLNK(status.st_mode):
			raise GameImportError(f"game installation contains a symbolic link: {relative}")
		if not stat.S_ISREG(status.st_mode):
			raise GameImportError(f"game installation contains a non-file: {relative}")

	files = 0
	total_bytes = 0

	def walk_error(error: OSError) -> None:
		raise GameImportError(f"cannot inventory game tree: {error}") from error

	for root, directories, filenames in os.walk(
		source,
		followlinks=False,
		onerror=walk_error,
	):
		root_path = Path(root)
		for directory in directories:
			path = root_path / directory
			try:
				status = path.lstat()
			except OSError as exc:
				raise GameImportError(f"cannot inventory game directory {path}: {exc}") from exc
			if stat.S_ISLNK(status.st_mode):
				raise GameImportError(f"game installation contains a symbolic link: {path}")
			if not stat.S_ISDIR(status.st_mode):
				raise GameImportError(f"game installation contains a special entry: {path}")
		for filename in filenames:
			path = root_path / filename
			try:
				status = path.lstat()
			except OSError as exc:
				raise GameImportError(f"cannot inventory game file {path}: {exc}") from exc
			if stat.S_ISLNK(status.st_mode):
				raise GameImportError(f"game installation contains a symbolic link: {path}")
			if not stat.S_ISREG(status.st_mode):
				raise GameImportError(f"game installation contains a special entry: {path}")
			total_bytes += status.st_size
			files += 1
	return GameTreeInventory(
		files=files,
		bytes=total_bytes,
		core_sha256={str(relative): _sha256(source / relative) for relative in REQUIRED_GAME_FILES},
	)


def discover_existing_game(
	paths: AppPaths,
	settings: Settings,
	environment: Mapping[str, str] | None = None,
) -> Path | None:
	env = os.environ if environment is None else environment
	candidates: list[Path] = []
	source_override = settings.source_override
	if source_override is not None:
		candidates.append(source_override)
	home_value = env.get("HOME")
	if home_value:
		home = Path(home_value).expanduser()
		for root_name in ("PristonTale-EU", "pristontale-eu"):
			root = home / "Games" / root_name
			candidates.extend(
				(
					root / "wine/drive_c/Program Files (x86)/Pristontale EU",
					root / "drive_c/Program Files (x86)/Pristontale EU",
					root / "Pristontale EU",
				)
			)

	destination = paths.source_game.resolve(strict=False)
	seen: set[Path] = set()
	for candidate in candidates:
		resolved = candidate.expanduser().resolve(strict=False)
		if resolved in seen or resolved == destination:
			continue
		seen.add(resolved)
		try:
			inventory_game_tree(resolved)
		except (FileNotFoundError, GameImportError, OSError):
			continue
		return resolved
	return None


def _validate_destination(paths: AppPaths, destination: Path) -> None:
	expected = (
		paths.launcher_prefix
		/ "drive_c/Program Files (x86)/Pristontale EU"
	).resolve(strict=False)
	if destination.resolve(strict=False) != expected:
		raise GameImportError(f"unexpected managed game destination: {destination}")


def _install_depot_metadata(
	paths: AppPaths,
	metadata: DepotMetadata,
	environment: Mapping[str, str] | None,
) -> Path:
	env = os.environ if environment is None else environment
	users_root = paths.launcher_prefix / "drive_c/users"
	preferred_names = [env.get("USER"), Path(env.get("HOME", "")).name]
	user_home = None
	for name in preferred_names:
		if not name:
			continue
		candidate = users_root / name
		if candidate.is_dir():
			user_home = candidate
			break
	if user_home is None:
		name = next((value for value in preferred_names if value), None)
		if name is None:
			raise GameImportError("cannot determine the Wine user for depot metadata")
		user_home = users_root / name
	target = user_home / "AppData/Roaming/Pristontale EU Launcher/Depots/depot_1.json"
	target.parent.mkdir(parents=True, exist_ok=True)
	temporary = target.with_name(f".{target.name}.{os.getpid()}.tmp")
	try:
		shutil.copy2(metadata.path, temporary)
		if _sha256(temporary) != metadata.sha256:
			raise GameImportError("copied depot metadata hash mismatch")
		temporary.replace(target)
	except BaseException:
		temporary.unlink(missing_ok=True)
		raise
	return target


def _remove_staging_tree(path: Path, expected_parent: Path) -> None:
	if path.parent.resolve() != expected_parent.resolve():
		raise GameImportError(f"refusing to remove staging path outside destination: {path}")
	if path.is_symlink() or path.is_file():
		path.unlink(missing_ok=True)
	elif path.is_dir():
		shutil.rmtree(path)


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


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()}.{uuid.uuid4().hex}.tmp")
	try:
		with temporary.open("x", encoding="utf-8") 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


def _write_import_manifest(paths: AppPaths, result: GameImportResult) -> None:
	_write_json(
		paths.game_import_manifest,
		{"schema": 2, **asdict(result)},
	)