sc2process.py

from __future__ import annotations

import asyncio
import os
import os.path
import shutil
import signal
import subprocess
import sys
import tempfile
import time
from contextlib import suppress
from pathlib import Path
from typing import Any

import aiohttp

# pyre-ignore[21]
import portpicker
from aiohttp.client_ws import ClientWebSocketResponse
from loguru import logger

from sc2 import paths, wsl
from sc2.controller import Controller
from sc2.paths import Paths
from sc2.versions import VERSIONS


class KillSwitch:
    _to_kill: list[Any] = []

    @classmethod
    def add(cls, value) -> None:
        logger.debug("kill_switch: Add switch")
        cls._to_kill.append(value)

    @classmethod
    def kill_all(cls) -> None:
        logger.info(f"kill_switch: Process cleanup for {len(cls._to_kill)} processes")
        for p in cls._to_kill:
            p._clean(verbose=False)


class SC2Process:
    """
    A class for handling SCII applications.

    :param host: hostname for the url the SCII application will listen to
    :param port: the websocket port the SCII application will listen to
    :param fullscreen: whether to launch the SCII application in fullscreen or not, defaults to False
    :param resolution: (window width, window height) in pixels, defaults to (1024, 768)
    :param placement: (x, y) the distances of the SCII app's top left corner from the top left corner of the screen
                       e.g. (20, 30) is 20 to the right of the screen's left border, and 30 below the top border
    :param render:
    :param sc2_version:
    :param base_build:
    :param data_hash:
    """

    def __init__(
        self,
        host: str | None = None,
        port: int | None = None,
        fullscreen: bool = False,
        resolution: list[int] | tuple[int, int] | None = None,
        placement: list[int] | tuple[int, int] | None = None,
        render: bool = False,
        sc2_version: str | None = None,
        base_build: str | None = None,
        data_hash: str | None = None,
    ) -> None:
        assert isinstance(host, str) or host is None
        assert isinstance(port, int) or port is None

        self._render = render
        self._arguments: dict[str, str] = {"-displayMode": str(int(fullscreen))}
        if not fullscreen:
            if resolution and len(resolution) == 2:
                self._arguments["-windowwidth"] = str(resolution[0])
                self._arguments["-windowheight"] = str(resolution[1])
            if placement and len(placement) == 2:
                self._arguments["-windowx"] = str(placement[0])
                self._arguments["-windowy"] = str(placement[1])

        self._host = host or os.environ.get("SC2CLIENTHOST", "127.0.0.1")
        self._serverhost = os.environ.get("SC2SERVERHOST", self._host)

        if port is None:
            self._port = portpicker.pick_unused_port()
        else:
            self._port = port
        self._used_portpicker = bool(port is None)
        self._tmp_dir = tempfile.mkdtemp(prefix="SC2_")
        self._process: subprocess.Popen | None = None
        self._session = None
        self._ws = None
        self._sc2_version = sc2_version
        self._base_build = base_build
        self._data_hash = data_hash

    async def __aenter__(self) -> Controller:
        KillSwitch.add(self)

        def signal_handler(*_args):
            # unused arguments: signal handling library expects all signal
            # callback handlers to accept two positional arguments
            KillSwitch.kill_all()

        signal.signal(signal.SIGINT, signal_handler)

        try:
            self._process = self._launch()
            self._ws = await self._connect()
        except:
            await self._close_connection()
            self._clean()
            raise

        return Controller(self._ws, self)

    async def __aexit__(self, *args) -> None:
        await self._close_connection()
        KillSwitch.kill_all()
        signal.signal(signal.SIGINT, signal.SIG_DFL)

    @property
    def ws_url(self) -> str:
        return f"ws://{self._host}:{self._port}/sc2api"

    @property
    def versions(self):
        """Opens the versions.json file which origins from
        https://github.com/Blizzard/s2client-proto/blob/master/buildinfo/versions.json"""
        return VERSIONS

    def find_data_hash(self, target_sc2_version: str) -> str | None:
        """Returns the data hash from the matching version string."""
        version: dict
        for version in self.versions:
            if version["label"] == target_sc2_version:
                return version["data-hash"]
        return None

    def _launch(self):
        if self._base_build:
            executable = str(paths.latest_executeble(Paths.BASE / "Versions", self._base_build))
        else:
            executable = str(Paths.EXECUTABLE)
        if self._port is None:
            self._port = portpicker.pick_unused_port()
            self._used_portpicker = True
        args = paths.get_runner_args(Paths.CWD) + [
            executable,
            "-listen",
            self._serverhost,
            "-port",
            str(self._port),
            "-dataDir",
            str(Paths.BASE),
            "-tempDir",
            self._tmp_dir,
        ]
        for arg, value in self._arguments.items():
            args.append(arg)
            args.append(value)
        if self._sc2_version:

            def special_match(strg: str):
                """Tests if the specified version is in the versions.py dict."""
                return any(version["label"] == strg for version in self.versions)

            valid_version_string = special_match(self._sc2_version)
            if valid_version_string:
                self._data_hash = self.find_data_hash(self._sc2_version)
                assert (
                    self._data_hash is not None
                ), f"StarCraft 2 Client version ({self._sc2_version}) was not found inside sc2/versions.py file. Please check your spelling or check the versions.py file."

            else:
                logger.warning(
                    f'The submitted version string in sc2.rungame() function call (sc2_version="{self._sc2_version}") was not found in versions.py. Running latest version instead.'
                )

        if self._data_hash:
            args.extend(["-dataVersion", self._data_hash])

        if self._render:
            args.extend(["-eglpath", "libEGL.so"])

        # if logger.getEffectiveLevel() <= logging.DEBUG:
        args.append("-verbose")

        sc2_cwd = str(Paths.CWD) if Paths.CWD else None

        if paths.PF in {"WSL1", "WSL2"}:
            return wsl.run(args, sc2_cwd)

        return subprocess.Popen(
            args,
            cwd=sc2_cwd,
            # Suppress Wine error messages
            stderr=subprocess.DEVNULL,
            # , env=run_config.env
        )

    async def _connect(self) -> ClientWebSocketResponse:
        # How long it waits for SC2 to start (in seconds)
        for i in range(180):
            if self._process is None:
                # The ._clean() was called, clearing the process
                logger.debug("Process cleanup complete, exit")
                sys.exit()

            await asyncio.sleep(1)
            try:
                self._session = aiohttp.ClientSession()
                ws = await self._session.ws_connect(self.ws_url, timeout=120)
                # FIXME fix deprecation warning in for future aiohttp version
                # ws = await self._session.ws_connect(
                #     self.ws_url, timeout=aiohttp.client_ws.ClientWSTimeout(ws_close=120)
                # )
                logger.debug("Websocket connection ready")
                return ws
            except aiohttp.client_exceptions.ClientConnectorError:
                await self._session.close()
                if i > 15:
                    logger.debug("Connection refused (startup not complete (yet))")

        logger.debug("Websocket connection to SC2 process timed out")
        raise TimeoutError("Websocket")

    async def _close_connection(self) -> None:
        logger.info(f"Closing connection at {self._port}...")

        if self._ws is not None:
            await self._ws.close()

        if self._session is not None:
            await self._session.close()

    def _clean(self, verbose: bool = True) -> None:
        if verbose:
            logger.info("Cleaning up...")

        if self._process is not None:
            assert isinstance(self._process, subprocess.Popen)
            if paths.PF in {"WSL1", "WSL2"}:
                if wsl.kill(self._process):
                    logger.error("KILLED")
            elif self._process.poll() is None:
                for _ in range(3):
                    self._process.terminate()
                    time.sleep(0.5)
                    if not self._process or self._process.poll() is not None:
                        break
            else:
                self._process.kill()
                self._process.wait()
                logger.error("KILLED")
            # Try to kill wineserver on linux
            if paths.PF in {"Linux", "WineLinux"}:
                # Command wineserver not detected
                with suppress(FileNotFoundError), subprocess.Popen(["wineserver", "-k"]) as p:
                    p.wait()

        if Path(self._tmp_dir).exists():
            shutil.rmtree(self._tmp_dir)

        self._process = None
        self._ws = None
        if self._used_portpicker and self._port is not None:
            portpicker.return_port(self._port)
            self._port = None
        if verbose:
            logger.info("Cleanup complete")