pristontale-eu-linux/ptlinux/login_trace.py

"""Privacy-safe instrumentation for the post-authentication login transition."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from datetime import datetime
import hashlib
import json
import os
from pathlib import Path
import platform
import re
import tempfile
from typing import Iterable

from .command import CommandOutcome, CommandRunner
from .login_diagnostics import latest_native_game_log
from .paths import AppPaths
from .runtime import (
	RuntimeSelection,
	detect_distribution_id,
	resolve_proton_sync_mode,
	run_game,
	run_proton_prefix_command,
)
from .settings import Settings


RELAY_REGISTRY_KEY = r"HKCU\Software\Wine\Debug"
RELAY_REGISTRY_VALUE = "RelayInclude"
RELAY_INCLUDE_FILTER = "KERNEL32.lstrcmpA"
GAME_DLL_LOGIN_COMPARE_RETURN_OFFSET = 0x79F157
DEFAULT_TRACE_TIMEOUT_SECONDS = 10 * 60.0

_LOAD_GAME_DLL = re.compile(
	r'Loaded L"[^"\r\n]*(?:\\|/)game\.dll" at (?P<base>[0-9a-f]+):',
	re.IGNORECASE,
)
_COMPARE_RETURN = re.compile(
	r'(?P<thread>[0-9a-f]{4}):Ret\s+KERNEL32\.lstrcmpA\(\) '
	r'retval=(?P<result>[0-9a-f]+) ret=(?P<return_address>[0-9a-f]+)',
	re.IGNORECASE,
)
_RELAY_VALUE = re.compile(
	rf"^\s*{re.escape(RELAY_REGISTRY_VALUE)}\s+REG_SZ(?:\s+(?P<value>.*))?$",
	re.MULTILINE,
)
_RELAY_CALL = re.compile(
	r'^(?P<prefix>.*?Call\s+KERNEL32\.lstrcmpA)\(.*\) '
	r'ret=(?P<return_address>[0-9a-f]+)\s*$',
	re.IGNORECASE,
)


@dataclass(frozen=True)
class RelayTraceResult:
	game_dll_base: int | None
	target_return_address: int | None
	comparison_results: tuple[bool, ...]
	log_files: int


@dataclass(frozen=True)
class LoginTraceReport:
	started_at: str
	finished_at: str
	result_directory: str
	summary_file: str
	report_file: str
	distribution_id: str
	kernel_release: str
	runtime_kind: str
	runtime_path: str | None
	proton_sync_effective: str | None
	game_exe_sha256: str | None
	game_dll_sha256: str | None
	native_log_before: str | None
	native_log_before_modified_at: str | None
	native_log_after: str | None
	native_log_after_modified_at: str | None
	game_returncode: int | None
	game_timed_out: bool
	game_duration_seconds: float
	registry_restored: bool
	raw_relay_deleted: bool
	relay_log_files: int
	game_dll_base: str | None
	target_return_address: str | None
	comparison_observed: bool
	account_id_exact_match: bool | None
	comparison_attempts: int
	conclusion: str
	warnings: tuple[str, ...]

	def to_dict(self) -> dict[str, object]:
		return asdict(self)


@dataclass(frozen=True)
class _RegistryValue:
	existed: bool
	value: str | None


def run_login_trace(
	paths: AppPaths,
	settings: Settings,
	selection: RuntimeSelection,
	runner: CommandRunner,
	*,
	use_gamebox: bool = True,
	timeout_seconds: float = DEFAULT_TRACE_TIMEOUT_SECONDS,
) -> LoginTraceReport:
	"""Launch one instrumented login attempt and retain only a redacted summary."""

	if timeout_seconds <= 0:
		raise ValueError("login trace timeout must be greater than zero")
	if selection.kind not in {"proton7", "proton7-container"}:
		raise RuntimeError("login tracing currently requires the pinned Proton 7 runtime")
	if not selection.proton_path:
		raise RuntimeError("the selected Proton 7 runtime path is incomplete")

	started = datetime.now().astimezone()
	result_directory = _create_result_directory(paths, started)
	game_exe = paths.runtime_view / "Game.exe"
	game_dll = paths.runtime_view / "game.dll"
	log_before, log_before_modified = _native_log_snapshot(paths, settings)
	registry_before = _query_relay_filter(paths, selection, runner)
	registry_restored = False
	trace = RelayTraceResult(None, None, (), 0)
	outcome: CommandOutcome | None = None
	raw_directory_path: Path | None = None

	paths.cache.mkdir(parents=True, exist_ok=True)
	with tempfile.TemporaryDirectory(
		prefix="ptlinux-login-relay-",
		dir=paths.cache,
	) as raw_directory:
		raw_directory_path = Path(raw_directory)
		try:
			_set_relay_filter(paths, selection, runner, RELAY_INCLUDE_FILTER)
			outcome = run_game(
				paths,
				settings,
				selection,
				runner,
				use_gamebox=use_gamebox,
				use_world_fix=False,
				timeout_seconds=timeout_seconds,
				check=False,
				extra_environment={
					"PROTON_LOG": "1",
					"PROTON_LOG_DIR": str(raw_directory_path),
					"SteamAppId": "0",
					"SteamGameId": "0",
					"WINEDEBUG": "+relay,+loaddll",
				},
				output_line_filter=redact_relay_output_line,
			)
			trace = parse_relay_logs(raw_directory_path.rglob("*.log"))
		finally:
			registry_restored = _restore_relay_filter(
				paths,
				selection,
				runner,
				registry_before,
			)

	if outcome is None:
		raise RuntimeError("the instrumented game process did not start")
	raw_relay_deleted = raw_directory_path is not None and not raw_directory_path.exists()
	log_after, log_after_modified = _native_log_snapshot(paths, settings)
	comparison_observed = bool(trace.comparison_results)
	account_id_exact_match = (
		trace.comparison_results[-1] if comparison_observed else None
	)
	conclusion, warnings = _conclusion(
		trace,
		registry_restored=registry_restored,
		raw_relay_deleted=raw_relay_deleted,
		timed_out=outcome.timed_out,
	)
	finished = datetime.now().astimezone()
	summary_file = result_directory / "summary.json"
	report_file = result_directory / "report.md"
	report = LoginTraceReport(
		started_at=started.isoformat(),
		finished_at=finished.isoformat(),
		result_directory=str(result_directory),
		summary_file=str(summary_file),
		report_file=str(report_file),
		distribution_id=detect_distribution_id() or "unknown",
		kernel_release=platform.release(),
		runtime_kind=selection.kind,
		runtime_path=selection.proton_path,
		proton_sync_effective=resolve_proton_sync_mode(settings.proton_sync),
		game_exe_sha256=_sha256_if_file(game_exe),
		game_dll_sha256=_sha256_if_file(game_dll),
		native_log_before=log_before,
		native_log_before_modified_at=log_before_modified,
		native_log_after=log_after,
		native_log_after_modified_at=log_after_modified,
		game_returncode=outcome.returncode,
		game_timed_out=outcome.timed_out,
		game_duration_seconds=round(outcome.duration_seconds, 3),
		registry_restored=registry_restored,
		raw_relay_deleted=raw_relay_deleted,
		relay_log_files=trace.log_files,
		game_dll_base=_hex_or_none(trace.game_dll_base),
		target_return_address=_hex_or_none(trace.target_return_address),
		comparison_observed=comparison_observed,
		account_id_exact_match=account_id_exact_match,
		comparison_attempts=len(trace.comparison_results),
		conclusion=conclusion,
		warnings=warnings,
	)
	_write_report(report)
	return report


def parse_relay_logs(log_paths: Iterable[Path]) -> RelayTraceResult:
	"""Extract only module addresses and equality results from raw Proton logs."""

	game_dll_base: int | None = None
	returns: list[tuple[int, int]] = []
	log_files = 0
	for path in log_paths:
		try:
			handle = path.open("r", encoding="utf-8", errors="replace")
		except OSError:
			continue
		log_files += 1
		with handle:
			for line in handle:
				if game_dll_base is None:
					loaded = _LOAD_GAME_DLL.search(line)
					if loaded is not None:
						game_dll_base = int(loaded.group("base"), 16)
				returned = _COMPARE_RETURN.search(line)
				if returned is not None:
					returns.append(
						(
							int(returned.group("return_address"), 16),
							int(returned.group("result"), 16),
						)
					)
	target_return_address = (
		game_dll_base + GAME_DLL_LOGIN_COMPARE_RETURN_OFFSET
		if game_dll_base is not None
		else None
	)
	comparison_results = tuple(
		result == 0
		for return_address, result in returns
		if target_return_address is not None and return_address == target_return_address
	)
	return RelayTraceResult(
		game_dll_base=game_dll_base,
		target_return_address=target_return_address,
		comparison_results=comparison_results,
		log_files=log_files,
	)


def redact_relay_output_line(line: str) -> str | None:
	"""Remove strings from a relay call before console or command-log output."""

	match = _RELAY_CALL.match(line.rstrip("\r\n"))
	if match is None:
		return line
	return (
		f"{match.group('prefix')}(<redacted>) "
		f"ret={match.group('return_address')}\n"
	)


def _query_relay_filter(
	paths: AppPaths,
	selection: RuntimeSelection,
	runner: CommandRunner,
) -> _RegistryValue:
	outcome = run_proton_prefix_command(
		paths,
		selection,
		runner,
		[
			"reg.exe",
			"query",
			RELAY_REGISTRY_KEY,
			"/v",
			RELAY_REGISTRY_VALUE,
		],
		check=False,
	)
	match = _RELAY_VALUE.search(outcome.stdout)
	if match is None:
		return _RegistryValue(False, None)
	return _RegistryValue(True, match.group("value") or "")


def _set_relay_filter(
	paths: AppPaths,
	selection: RuntimeSelection,
	runner: CommandRunner,
	value: str,
) -> None:
	run_proton_prefix_command(
		paths,
		selection,
		runner,
		[
			"reg.exe",
			"add",
			RELAY_REGISTRY_KEY,
			"/v",
			RELAY_REGISTRY_VALUE,
			"/t",
			"REG_SZ",
			"/d",
			value,
			"/f",
		],
	)


def _restore_relay_filter(
	paths: AppPaths,
	selection: RuntimeSelection,
	runner: CommandRunner,
	previous: _RegistryValue,
) -> bool:
	if previous.existed:
		outcome = run_proton_prefix_command(
			paths,
			selection,
			runner,
			[
				"reg.exe",
				"add",
				RELAY_REGISTRY_KEY,
				"/v",
				RELAY_REGISTRY_VALUE,
				"/t",
				"REG_SZ",
				"/d",
				previous.value or "",
				"/f",
			],
			check=False,
		)
	else:
		outcome = run_proton_prefix_command(
			paths,
			selection,
			runner,
			[
				"reg.exe",
				"delete",
				RELAY_REGISTRY_KEY,
				"/v",
				RELAY_REGISTRY_VALUE,
				"/f",
			],
			check=False,
		)
	return not outcome.timed_out and outcome.returncode == 0


def _create_result_directory(paths: AppPaths, started: datetime) -> Path:
	root = paths.state / "login-traces"
	root.mkdir(parents=True, exist_ok=True)
	stem = f"{started.strftime('%Y%m%d-%H%M%S')}-{os.getpid()}"
	result = root / stem
	suffix = 1
	while result.exists():
		result = root / f"{stem}-{suffix}"
		suffix += 1
	result.mkdir(mode=0o700)
	return result


def _native_log_snapshot(
	paths: AppPaths,
	settings: Settings,
) -> tuple[str | None, str | None]:
	path = latest_native_game_log(paths, settings)
	if path is None:
		return None, None
	try:
		modified = datetime.fromtimestamp(path.stat().st_mtime).astimezone().isoformat()
	except OSError:
		modified = None
	return str(path), modified


def _conclusion(
	trace: RelayTraceResult,
	*,
	registry_restored: bool,
	raw_relay_deleted: bool,
	timed_out: bool,
) -> tuple[str, tuple[str, ...]]:
	warnings: list[str] = []
	if trace.comparison_results:
		if trace.comparison_results[-1]:
			conclusion = "account-id-match"
		else:
			conclusion = "account-id-mismatch"
			warnings.append(
				"The authenticated reply was discarded because its account ID did not "
				"exactly match the ID entered in the client."
			)
	else:
		conclusion = "comparison-not-observed"
		warnings.append(
			"The post-authentication account comparison was not observed; confirm that "
			"Login was pressed before closing the game."
		)
	if trace.log_files == 0:
		warnings.append("Proton did not create a relay log for this run.")
	elif trace.game_dll_base is None:
		warnings.append("The relay log did not record the loaded game.dll base address.")
	if timed_out:
		warnings.append("The diagnostic reached its timeout and terminated the launcher process.")
	if not registry_restored:
		warnings.append("The temporary Wine relay registry value could not be restored.")
	if not raw_relay_deleted:
		warnings.append("The private raw relay directory could not be removed.")
	return conclusion, tuple(warnings)


def _write_report(report: LoginTraceReport) -> None:
	summary_path = Path(report.summary_file)
	report_path = Path(report.report_file)
	summary_path.write_text(
		json.dumps(report.to_dict(), indent=2, sort_keys=True) + "\n",
		encoding="utf-8",
	)
	match = (
		"not observed"
		if report.account_id_exact_match is None
		else "yes" if report.account_id_exact_match else "no"
	)
	report_path.write_text(
		"# PristonTale EU login trace\n\n"
		f"- Started: {report.started_at}\n"
		f"- Runtime: {report.runtime_kind}\n"
		f"- Account-ID comparison observed: {report.comparison_observed}\n"
		f"- Exact account-ID match: {match}\n"
		f"- Conclusion: `{report.conclusion}`\n"
		f"- Registry restored: {report.registry_restored}\n"
		f"- Raw relay data deleted: {report.raw_relay_deleted}\n\n"
		"No account names, passwords, or packet payloads are retained in this report.\n",
		encoding="utf-8",
	)


def _sha256_if_file(path: Path) -> str | None:
	try:
		if not path.is_file():
			return None
		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()
	except OSError:
		return None


def _hex_or_none(value: int | None) -> str | None:
	return f"0x{value:08x}" if value is not None else None