"""Experiment cleanup, prefix-link rollback, and registry restoration.""" from __future__ import annotations from collections.abc import Callable, Mapping, Sequence import os from pathlib import Path from ..config import AppConfig from ..prefix import cleanup_dedicated_wineserver from ..registry import DesktopRegistrySnapshot from ..result import ResultDirectory from ..runner import CommandLogger, CommandResult from .models import ExperimentError EXPLORER_KEY = r"HKCU\Software\Wine\Explorer" DESKTOPS_KEY = r"HKCU\Software\Wine\Explorer\Desktops" DISPOSABLE_DIRECTORY_NAME = "ptdiag-disposable" ProtonCommand = Callable[..., list[str]] WineserverCleaner = Callable[..., CommandResult] def replace_game_link( prefix_link: Path, disposable_game: Path, expected_live_game: Path, logger: CommandLogger, ) -> str: """Point the dedicated prefix at a verified disposable game tree.""" if not prefix_link.is_symlink(): raise ExperimentError( f"Dedicated-prefix game path is not a symlink: {prefix_link}" ) if prefix_link.resolve(strict=True) != expected_live_game.resolve(strict=True): raise ExperimentError("Dedicated-prefix game link changed after source preflight") original_target = os.readlink(prefix_link) require_command( logger.run( ["ln", "-sfn", "--", str(disposable_game), str(prefix_link)], cwd=prefix_link.parent, timeout_seconds=15.0, ), "point dedicated prefix at disposable game", ) if ( not prefix_link.is_symlink() or prefix_link.resolve(strict=True) != disposable_game.resolve(strict=True) ): raise ExperimentError( "Dedicated-prefix game link did not resolve to disposable copy" ) return original_target def restore_game_link( prefix_link: Path, original_target: str, expected_live_game: Path, logger: CommandLogger, ) -> None: """Restore the dedicated-prefix game link and verify its exact target.""" require_command( logger.run( ["ln", "-sfn", "--", original_target, str(prefix_link)], cwd=prefix_link.parent, timeout_seconds=15.0, ), "restore dedicated-prefix game link", ) if not prefix_link.is_symlink(): raise ExperimentError("Restored dedicated-prefix game path is not a symlink") if os.readlink(prefix_link) != original_target: raise ExperimentError("Dedicated-prefix game link text was not restored exactly") if prefix_link.resolve(strict=True) != expected_live_game.resolve(strict=True): raise ExperimentError("Dedicated-prefix game link target was not restored") def remove_disposable_game( config: AppConfig, result: ResultDirectory, disposable_game: Path, logger: CommandLogger, *, directory_name: str = DISPOSABLE_DIRECTORY_NAME, ) -> None: """Remove only the result-scoped disposable tree beneath its managed root.""" staging_root = (config.paths.app_root / directory_name).absolute() validate_disposable_path( staging_root, disposable_game, result.path.name, directory_name=directory_name, ) require_command( logger.run( ["rm", "-rf", "--one-file-system", "--", str(disposable_game)], timeout_seconds=900.0, ), "remove disposable game tree", ) if disposable_game.exists() or disposable_game.is_symlink(): raise ExperimentError( f"Disposable game cleanup left its target: {disposable_game}" ) try: staging_root.rmdir() except OSError: pass def validate_disposable_path( staging_root: Path, path: Path, expected_name: str, *, directory_name: str = DISPOSABLE_DIRECTORY_NAME, ) -> None: """Reject cleanup targets outside the managed, result-scoped staging root.""" if path.parent != staging_root or path.name != expected_name: raise ExperimentError(f"Unsafe disposable game path: {path}") if staging_root.name != directory_name: raise ExperimentError(f"Unsafe disposable staging root: {staging_root}") def require_command(outcome: CommandResult, operation: str) -> None: if not command_ok(outcome): raise ExperimentError( f"Could not {operation} ({outcome.returncode}): " f"{outcome.error or outcome.stderr.strip() or 'command failed'}" ) def registry_setup_commands( config: AppConfig, desktop_name: str, desktop_size: str, *, proton_command: ProtonCommand, ) -> list[list[str]]: return [ proton_command( config, "runinprefix", "reg.exe", "add", EXPLORER_KEY, "/v", "Desktop", "/t", "REG_SZ", "/d", desktop_name, "/f", ), proton_command( config, "runinprefix", "reg.exe", "add", DESKTOPS_KEY, "/v", desktop_name, "/t", "REG_SZ", "/d", desktop_size, "/f", ), ] def registry_restore_commands( config: AppConfig, snapshot: DesktopRegistrySnapshot, *, proton_command: ProtonCommand, ) -> list[list[str]]: return [ registry_restore_command( config, EXPLORER_KEY, "Desktop", snapshot.selected_desktop.exists, snapshot.selected_desktop.value, proton_command=proton_command, ), registry_restore_command( config, DESKTOPS_KEY, snapshot.desktop_name, snapshot.desktop_size.exists, snapshot.desktop_size.value, proton_command=proton_command, ), ] def registry_restore_command( config: AppConfig, key: str, name: str, existed: bool, value: str | None, *, proton_command: ProtonCommand, ) -> list[str]: if existed: assert value is not None return proton_command( config, "runinprefix", "reg.exe", "add", key, "/v", name, "/t", "REG_SZ", "/d", value, "/f", ) return proton_command( config, "runinprefix", "reg.exe", "delete", key, "/v", name, "/f", ) def run_checked_commands( logger: CommandLogger, commands: Sequence[Sequence[str]], *, cwd: Path, env_overrides: Mapping[str, str], ) -> None: for command in commands: outcome = logger.run( command, cwd=cwd, env_overrides=env_overrides, timeout_seconds=15.0, ) if not command_ok(outcome): raise ExperimentError( f"Command failed ({outcome.returncode}): {list(command)!r}: " f"{outcome.error or outcome.stderr.strip()}" ) def cleanup_checked( config: AppConfig, logger: CommandLogger, *, cleaner: WineserverCleaner = cleanup_dedicated_wineserver, ) -> None: """Stop the dedicated wineserver while accepting its idle return contract.""" outcome = cleaner(config, logger, timeout_seconds=15.0) if ( outcome.returncode == 1 and not outcome.timed_out and outcome.error is None and not outcome.stdout and not outcome.stderr ): return if not command_ok(outcome): raise ExperimentError( "Dedicated-prefix wineserver cleanup failed: " f"{outcome.error or outcome.stderr.strip() or outcome.returncode}" ) def command_ok(result: CommandResult) -> bool: return not result.timed_out and result.error is None and result.returncode == 0